Compare commits
39
Commits
d1f780e381
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2832566e87 | ||
|
|
7d5d7b9e36 | ||
|
|
4497d80a94 | ||
|
|
04f1375b45 | ||
|
|
d9b2359d25 | ||
|
|
cc72248e26 | ||
|
|
c4154d4db9 | ||
|
|
0ed03ac095 | ||
|
|
2ea3fa14d8 | ||
|
|
dbfb925f65 | ||
|
|
67acae708c | ||
|
|
7e44025671 | ||
|
|
2dc79f0637 | ||
|
|
5894e64287 | ||
|
|
3e860c16e0 | ||
|
|
68a8e49896 | ||
|
|
5cb35945a4 | ||
|
|
f3bb09e6b5 | ||
|
|
6f1f2bd296 | ||
|
|
38be729509 | ||
|
|
d4e627fbb2 | ||
|
|
3e6e171831 | ||
|
|
d1357ae483 | ||
|
|
9db2bc112d | ||
|
|
9064750f33 | ||
|
|
5dc5ec9af2 | ||
|
|
e43aeeaf44 | ||
|
|
7dc6796f3a | ||
|
|
f7fc06c0ff | ||
|
|
04e3357dd0 | ||
|
|
2c6d1e8ce6 | ||
|
|
457bb900fb | ||
|
|
f4fcf175ec | ||
|
|
d513c5f892 | ||
|
|
4cb97318b5 | ||
|
|
31b2a4d6ca | ||
|
|
db54921066 | ||
|
|
9fee9a1478 | ||
|
|
d3acccd146 |
No files matched your search
@@ -28,3 +28,6 @@ hailort.log
|
||||
|
||||
# 录制产物(可复现:mtc DSL → 编译 → pybot --record)
|
||||
actpred/data_v3
|
||||
|
||||
# pybot 训练/评估数据(pybot_record_* 动作录制 + eval_recorded*.json)
|
||||
data/
|
||||
@@ -235,7 +235,8 @@
|
||||
|
||||
- 一套相对独立的单智能体产品,基于 DeepSeek + WebSocket(:8765),由 `core` 拉起但与 waf / gateway / dnfs2 等主干微服务无调用关系。
|
||||
- 自带技能系统(`skills/`)、任务调度器、记忆管理、前端页面与语音客户端(`voice_client.py`)。
|
||||
- 默认模型 `deepseek-v4-flash`,配置在 `config.json`;`models` 配置节可挂多个 OpenAI 兼容端点按档案名切换(档案级 reasoning/温度/展示名/`fallback` 失败转移),`--once` 支持无 WebSocket 的 CLI 一次性对话;详见 [aibox/README.md](aibox/README.md)。
|
||||
- 默认模型 `deepseek-v4-flash`,配置在 `config.json`;`models` 配置节可挂多个 OpenAI 兼容端点按档案名切换(档案级 reasoning/温度/展示名/`fallback` 失败转移),`model_config.py` 提供交互式接入点管理(列表/添加/删除/连通性验证,Key 仅记环境变量名不落盘),`--once` 支持无 WebSocket 的 CLI 一次性对话;详见 [aibox/README.md](aibox/README.md)。
|
||||
- **六轮安全审计(2026-08,Batch E/F/G/H/I/K/L)**:命令执行保护链(1MB/超时/中断)、凭据脱敏(代理密码不落盘、日志打码、Bearer 脱敏)、取消传播(子代理/情绪子代理/调度器超时)、DB 事务与 FTS5 转义、POSIX 进程组终止、任务系统入参校验(畸形时间/数字字段防调度器热循环)、前端 XSS 净化(DOMPurify + SRI)等;第六轮交叉验证 Batch K 零回归,R1~R3 低危已随 Batch L 落地(72 项发现全部修复);单测 490 项全过(Linux);审计报告见 [docs/reports/aibox_audit_20260817.md](docs/reports/aibox_audit_20260817.md)。
|
||||
|
||||
### BotOps(aibox 智能体 TUI 前端)
|
||||
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
"""多头 one-hot 分箱 —— 18 维连续动作向量与多组 one-hot 分类目标互转
|
||||
|
||||
分箱设计(actpred_onehot_multihead_dev_plan_20260814.md):
|
||||
分箱设计:
|
||||
0-9 舵机角:均匀 1°/箱 → 180 箱(匹配 0.75° 微动作插值数据粒度)
|
||||
10 spk.F:对数分箱 log10(1+x) ∈ [0, log10(5001)] → 32 箱
|
||||
11 spk.delay:对数分箱 log10(1+x) ∈ [0, log10(2001)] → 32 箱
|
||||
|
||||
@@ -2,8 +2,7 @@
|
||||
|
||||
从 .mt 文件中提取 @Action 块,解析为 (n_rows, n_cols) 的 float32 矩阵。
|
||||
支持 `-` 占位符(沿用上一行同列的值);首行 `-` 填 0,
|
||||
与 pybot/bot4cj 运行时语义一致(2026-08 统一,见
|
||||
docs/reports/bot4cj_python_migration_dev_doc.md §2.3)。
|
||||
与 pybot/bot4cj 运行时语义一致(2026-08 统一)。
|
||||
"""
|
||||
|
||||
import re
|
||||
|
||||
+99
-8
@@ -57,13 +57,28 @@ python aibox.py
|
||||
|
||||
启动后,WebSocket 服务器将运行在 `ws://localhost:8765`(地址与端口可在 `config.json` 的 `system.websocket_host` / `system.websocket_port` 中修改)。
|
||||
|
||||
> 注意:本模块入口文件为 `aibox.py`(当前约 3018 行,集成 WebSocket 服务器、AI 引擎、工具/技能系统、任务调度器于一体;自 v3.0 起工具类与子系统已抽离至 `tools/`、`subsystems/`、`common/`、`tests/` 包,aibox.py 仅保留核心类与编排逻辑)。`client.py` 为示例 WebSocket 客户端,`voice_client.py` 为语音客户端,`llm_client.py` 为多厂商大模型客户端抽象层,`_run.sh` 为被 `core` 拉起的启动脚本。
|
||||
### 管理大模型接入点(不用手写配置)
|
||||
|
||||
不想手写 `config.json` 的 `models` 节?用交互式菜单工具:
|
||||
|
||||
```bash
|
||||
python model_config.py # 交互式菜单:列出 / 添加 / 删除 / 连通性验证
|
||||
python model_config.py --list # 非交互:列出当前接入点
|
||||
```
|
||||
|
||||
- **添加接入点**:按提示填写档案名、API 地址(附常见厂商端点参考)、模型名;Key 两种方式——环境变量名(推荐,Key 不落盘)或免密本地端点(Ollama 等);可选显示名/温度/超时/思考模式/失败转移(从已有接入点多选)/设为默认
|
||||
- **删除接入点**:编号选择 + 确认,自动清理其他档案的失败转移引用
|
||||
- **连通性验证**:对选中接入点发一次最小请求,直接给出成败原因
|
||||
- 添加后可立即使用:`python aibox.py --model <档案名>` 指定,或设为默认后无需 `--model`
|
||||
|
||||
> 注意:本模块入口文件为 `aibox.py`(当前约 3408 行,集成 WebSocket 服务器、AI 引擎、工具/技能系统、任务调度器于一体;自 v3.0 起工具类与子系统已抽离至 `tools/`、`subsystems/`、`common/`、`tests/` 包,aibox.py 仅保留核心类与编排逻辑)。`client.py` 为示例 WebSocket 客户端,`voice_client.py` 为语音客户端,`llm_client.py` 为多厂商大模型客户端抽象层,`_run.sh` 为被 `core` 拉起的启动脚本。
|
||||
|
||||
## 📁 目录结构
|
||||
|
||||
```
|
||||
aibox/
|
||||
├── aibox.py # 服务主程序(WebSocket 服务器 + AI 引擎 + 任务调度器 + 核心类 DeepSeekChat/WebSocketHandler/EnhancedAIManager)
|
||||
├── model_config.py # 大模型接入点交互式管理工具(列表/添加/删除/连通性验证,不用手写配置)
|
||||
├── client.py # 示例 WebSocket 客户端
|
||||
├── voice_client.py # 语音客户端
|
||||
├── llm_client.py # 大模型客户端抽象层(多厂商切换,DeepSeekClient 兼容 OpenAI 协议)
|
||||
@@ -81,8 +96,9 @@ aibox/
|
||||
│ ├── db.py # MemoryDatabase + MemoDatabase(SQLite,FTS5 trigram 中文检索 + LIKE 兜底)
|
||||
│ ├── skill_manager.py # SkillManager(技能发现/加载/热重载)
|
||||
│ ├── conversation_history.py # ConversationHistory(对话历史管理)
|
||||
│ └── conversation_compressor.py # ConversationCompressor(对话压缩,注入 llm_client 解耦)
|
||||
├── tools/ # 工具系统包(A 批 + B 批共 17 个工具类)
|
||||
│ ├── conversation_compressor.py # ConversationCompressor(对话压缩,注入 llm_client 解耦)
|
||||
│ └── process_manager.py # ProcessManager(后台进程管理,全局单例 + 环形输出缓冲 + atexit 收尾)
|
||||
├── tools/ # 工具系统包(A 批 + B 批 + 后台进程共 22 个工具类)
|
||||
│ ├── base.py # Tool 抽象基类
|
||||
│ ├── registry.py # ToolRegistry 工具注册器
|
||||
│ ├── basic_tools.py # A 批:ReadFile/WriteFile/ListFiles/SimpleWebCrawler/GetCurrentTime/Calculator
|
||||
@@ -91,8 +107,9 @@ aibox/
|
||||
│ ├── db_tools.py # B 批:MemoryTool/MemoTool
|
||||
│ ├── prompt_tools.py # B 批:UpdateIdentity/UpdateSoul/SaveMemoryAndEndConversation
|
||||
│ ├── proxy_tools.py # B 批:SetProxy/GetProxy
|
||||
│ └── misc_tools.py # B 批:AskUserTool
|
||||
├── tests/ # 单元测试套件(355 项全过)
|
||||
│ ├── misc_tools.py # B 批:AskUserTool
|
||||
│ └── process_tools.py # 后台进程:BgRun/BgWrite/BgList/BgOutput/BgKill(bg_run/bg_write/bg_list/bg_output/bg_kill)
|
||||
├── tests/ # 单元测试套件(598 项全过)
|
||||
│ ├── conftest.py # sys.path 配置 + websockets 桩
|
||||
│ ├── test_basic_tools.py # A 批工具 43 项
|
||||
│ ├── test_b_tools.py # B 批工具 58 项
|
||||
@@ -103,7 +120,16 @@ aibox/
|
||||
│ ├── test_conversation_compressor.py # 对话压缩器
|
||||
│ ├── test_skill_manager.py # 技能管理器
|
||||
│ ├── test_once_mode.py # CLI 一次性对话模式(Batch A/B/C 15 项)
|
||||
│ └── test_models_profile.py # 多模型档案(Batch D 17 项)
|
||||
│ ├── test_models_profile.py # 多模型档案(Batch D 17 项)
|
||||
│ ├── test_config.py # 配置脱敏(F1 3 项)
|
||||
│ ├── test_f2_fixes.py # 超时清理/线程收尾(F2 5 项)
|
||||
│ ├── test_g_fixes.py # 第三轮审计 Batch G(20 项)
|
||||
│ ├── test_h_fixes.py # 第四轮审计 Batch H(26 项)
|
||||
│ ├── test_i_fixes.py # 遗留收尾 Batch I(16 项)
|
||||
│ ├── test_subagent_max_turns.py # 子代理最大轮次行为回归(14 项)
|
||||
│ ├── test_k_fixes.py # 第五轮审计 Batch K(35 项)
|
||||
│ ├── test_model_config.py # 模型接入点管理工具(17 项)
|
||||
│ └── test_process_manager.py # 后台进程管理器与 bg_* 工具(39 项:生命周期/交互 I/O/收尾)
|
||||
├── skills/ # 内置扩展技能(*.py 可执行脚本 / *.md Markdown 技能)
|
||||
│ ├── hello-world.py
|
||||
│ ├── aes-crypto.py
|
||||
@@ -138,7 +164,8 @@ common ← subsystems ← tools ← aibox.py
|
||||
```bash
|
||||
# 从仓库根目录运行
|
||||
py -m pytest aibox/tests -q
|
||||
# 当前:264 项全过
|
||||
# 当前:598 项全过(Linux;test_execute_timeout 为 Windows 向用例,本机无 ping 预存失败)
|
||||
# Windows 开发机:532 passed / 4 failed(平台环境项:symlink 权限、python3 缺失、POSIX 可执行位语义 ×2)/ 1 skipped
|
||||
```
|
||||
|
||||
运行期数据存放于 `~/.aibox/`:
|
||||
@@ -214,6 +241,13 @@ py -m pytest aibox/tests -q
|
||||
| `web` | 访问网页、获取标题/文本/链接(返回 markdown 指引,改用 `curl` 命令行访问) |
|
||||
| `calculator` | 数学计算 |
|
||||
| `get_current_time` | 获取当前时间 |
|
||||
| `bg_run` | 后台启动长驻进程(dev server / 训练任务 / watch),立即返回句柄不阻塞;stdin 为管道可交互;默认 `detach=false`(aibox 退出时终止),`detach=true` 豁免但脱离句柄管理;`log_file=true` 同步镜像全程输出到 `workspace/bg_<句柄>.log`(无界全档,可被 run_command 的 tail/cat 原生读取);名称参数为 `proc_name`(≤64 字符,防与 `execute_tool` 首参撞名) |
|
||||
| `bg_write` | 向后台进程 stdin 写入一行(行协议交互:REPL / 问答式 CLI);返回 **watermark 水位**——`bg_output(since=watermark)` 得到的必然且仅是本条输入引发的响应(写读对齐,不含 `[stdin]` 回显行);`wait_lines=N` 同步等待 N 行响应(expect 式一步问答);单次上限 64KB,超时写入不阻塞工具线程;无 TTY,全屏程序(vim/top)不可用,建议启动时强制行缓冲(`python -i -u` / `PYTHONUNBUFFERED=1` / `stdbuf -oL`) |
|
||||
| `bg_list` | 列出全部后台进程(运行中在前,已退出保留尾部输出供验尸,LRU 淘汰) |
|
||||
| `bg_output` | 拉取输出:`tail=N` 尾部 N 行 / `since=游标` 增量续读(游标为上次返回的 `next_cursor`);单次默认 64KB、上限 1MB;环形缓冲(默认 256KB/进程)淘汰的早期行如实提示 |
|
||||
| `bg_kill` | 终止后台进程(整棵进程组):三级优雅终止——关 stdin(行协议进程见 EOF 常自行退出)→ SIGTERM 宽限 → SIGKILL 兜底;`force=true` 直接 SIGKILL;对已退出进程幂等 |
|
||||
|
||||
> **后台进程语义(v3.8,T22 属主隔离)**:注册表为全局单例(跨会话/跨子代理可见——list/output 观察面全局共享;写/杀按属主隔离,仅启动者所在会话可 `bg_write`/`bg_kill` 自己的进程,防提示注入跨会话破坏他人 dev server;未标识属主的调用不受限);用户中断(Ctrl+C/取消事件)**不**影响后台进程;aibox 进程退出经 atexit 统一收尾(非 detach 全杀,SIGKILL/SIGTERM 硬杀不覆盖);单行输出截断 16KB;运行中上限 16 / 总条目上限 32 / 单次 stdin 写入上限 64KB(`background_processes` 配置节可调);POSIX 按进程组终止(`start_new_session`,孙进程不悬挂),Windows `taskkill /T /F` 收割进程树。
|
||||
|
||||
## 💬 WebSocket API
|
||||
|
||||
@@ -591,6 +625,14 @@ AI 可以通过 `update_soul` 工具动态更新 `99_SOUL.md`,通过 `update_i
|
||||
"commands": {
|
||||
"enabled": true
|
||||
},
|
||||
"background_processes": {
|
||||
"enabled": true,
|
||||
"max_processes": 16,
|
||||
"max_entries": 32,
|
||||
"buffer_per_process": 262144,
|
||||
"max_lifetime": 0,
|
||||
"max_write_bytes": 65536
|
||||
},
|
||||
"web": {
|
||||
"timeout": 10,
|
||||
"delay": 0,
|
||||
@@ -890,6 +932,55 @@ elif command == "/mycommand":
|
||||
|
||||
## 📄 更新日志
|
||||
|
||||
### v3.8
|
||||
- ✅ **后台进程池(智能体运行/管理长驻进程)**:新增 `subsystems/process_manager.py`(`ProcessManager` 全局单例)+ `tools/process_tools.py` 5 个 worker 执行面工具(`bg_run` / `bg_write` / `bg_list` / `bg_output` / `bg_kill`),弥补 `run_command` 同步阻塞、超时即杀的不足——dev server / 训练任务 / watch 等长驻进程可后台启动、跨会话管理:
|
||||
- **生命周期语义(三决策)**:用户中断(Ctrl+C/取消事件)不影响后台进程;aibox 退出经 atexit 统一收尾(非 detach 全杀,`detach=true` 豁免但脱离句柄管理);注册表跨会话/跨子代理共享(观察面全局可见;写/杀按属主隔离——T22,仅启动者所在会话可操作,定时任务 AI 为精简 `task` profile 且无 update_identity/update_soul/set_proxy)
|
||||
- **交互式 I/O(stdin 管道 + watermark 写读对齐)**:stdin 挂管道,`bg_write` 向行协议进程(REPL / 问答式 CLI)写入一行,返回 **watermark**(`[stdin]` 回显行序号)——`bg_output(since=watermark)` 得到的必然且仅是该输入引发的响应(写读配对);`wait_lines=N` expect 式同步等待;回显行入环形缓冲与 log 镜像(tail 可见完整交互时序);写入 watchdog 线程兜底(进程不读 stdin 管道写满时不阻塞工具线程);单次写入上限 64KB
|
||||
- **输出缓冲**:每进程有界环形缓冲(默认 256KB,`buffer_per_process` 可调),单行截断 16KB;`bg_output` 支持 `tail` 尾读与 `since` 游标增量续读(淘汰行如实提示),单次默认 64KB、上限 1MB;`bg_run log_file=true` 输出同步 tee 到 `workspace/bg_<句柄>.log`(无界全档,`[stdin]` 回显同入档,可被 run_command 的 tail/cat 原生读取)
|
||||
- **三级优雅终止**:`bg_kill` 关 stdin(行协议进程见 EOF 常自行退出)→ SIGTERM 宽限(grace 比例 4:6)→ SIGKILL 兜底;`force=true` 直接 SIGKILL;shutdown 收尾同样先关 stdin
|
||||
- **安全边界**:运行中上限 16 / 总条目上限 32(LRU 淘汰已退出条目);`background_processes` 配置节(含 `enabled` 双闸:`commands.enabled` 关闭时一并拒绝);`proc_name` 标签 ≤64 字符且拒绝控制字符(M5 惯例);POSIX `start_new_session` + 进程组终止(对齐执行器 L9 语义),Windows `taskkill /T /F` 收割进程树;TTL `max_lifetime` 惰性清理(默认 0 不启用);凭据类交互(sudo -S 喂密码)会留对话记录,敏感凭据走环境变量注入(工具描述已指引)
|
||||
- **实现要点**:H2 reader 线程模式复用(EOF 即 poll 收割防僵尸);`list_procs` 惰性刷新全表状态;`proc_name` 命名规避 `ToolRegistry.execute_tool(name=...)` 首参撞名;编排者误调拦截表 `WORKER_ONLY_TOOLS` 同步纳入;无 TTY 的适用边界(全屏程序不可用、行缓冲指引)写入工具描述供模型自检
|
||||
- **测试**:新增 `tests/test_process_manager.py` 39 项(生命周期/游标续读/环形淘汰/单行截断/上限/TTL/收尾/工具分发/禁用闸/watermark 对齐/回显/sync 等待/超时/EOF 优雅终止/log tee/注册表分发),Windows 开发机 532 passed / 4 failed(平台环境项,与代码无关)
|
||||
|
||||
### v3.7
|
||||
- ✅ **第六轮安全审计 + Batch L 修复(2026-08-20,交叉验证轮:Batch K 回归验证 / 并发与生命周期 / 降级路径 / 修复间相互作用,4 路并行审读 + 主审逐条复核)**(完整记录见 [docs/reports/aibox_audit_20260817.md](../docs/reports/aibox_audit_20260817.md) §16):
|
||||
- **Batch K 回归验证**:J1~J16 逐项在位、零回归,无修复引入的缺陷;历届 Batch E~I 修复抽查(executors reader 收尾、1MB 输出上限、registry 流式预算、broadcast 快照等)均在位
|
||||
- **Batch L(R1~R3 低危全落地)**:R1 `/memos` WS 命令 strptime 容错(畸形 reminder_time 按"待执行"展示,不再整体"处理失败");R2 移除 `handle_interrupt` 对全局 `ai_manager.command_executor` 的取消(WS 中断不再误杀 CLI 主 AI 子进程);R3 `unregister` 补齐回收(关闭 websocket + 取消在途处理任务,broadcast 超时路径与连接关闭路径行为对齐,半开连接不再泄漏)
|
||||
- **测试**:新增 `tests/test_l_fixes.py` 6 例(R1 容错与语义不回归 ×2 / R2 范围收窄与中断能力不回归 ×2 / R3 回收与幂等 ×2);Linux 维持 490 passed,Windows 开发机 493 passed / 4 failed(均为平台环境项,与代码无关)
|
||||
|
||||
### v3.6
|
||||
- ✅ **交互式大模型接入点管理([model_config.py](model_config.py))**:菜单式 列表 / 添加 / 删除 / 连通性验证,不用手写 `config.json`;Key 仅记录环境变量名不落盘(空 = 免密本地端点);删除自动清理其他档案的 fallback 引用;默认档案经 `api.default_model` 生效(详见 [模型接入点管理](#管理大模型接入点不用手写配置))
|
||||
- ✅ **第五轮安全审计(2026-08-19,4 路并行审读 + Linux PoC 2/2 实证,发现 J1~J16:高危 2 / 中危 5 / 低危 6 / 信息 3)**(完整记录见 [docs/reports/aibox_audit_20260817.md](../docs/reports/aibox_audit_20260817.md) §13~§14):
|
||||
- **K1(J1 高危·任务系统投毒)**:畸形时间/数字字段经 CLI/WS/memo 工具三入口零校验入库 → 调度器失败热循环 + 每个新连接握手即崩——入库端统一校验(`reminder_time`/`repeat_end_time` strptime 预检、`repeat_interval` 强转 int、`repeat_type` 白名单)、调度器预检失败按到期执行(Q5 语义)、`get_pending_tasks_for_user` 解析失败按到期展示、`register` 移入 try/finally(会话登记不再泄漏);附带修复 `add_user_task` 缺失 `repeat_end_time` 参数透传
|
||||
- **K2(J2 高危·前端 XSS)**:三页面 7 处 `marked.parse` 直渲(AI 输出/命令结果全链路 innerHTML)统一改经 `renderMarkdown()` + DOMPurify(缺失时退化为转义文本防 XSS 兜底);`entry1.addSystemMessage` 改 DOM 构建 + textContent;任务标题改 createTextNode;CDN 脚本(marked/hljs/socket.io/DOMPurify)补 SRI integrity(J13)
|
||||
- **K3(中危)**:FTS 外部内容表删除换序——先删 `memories_fts` 再删 `memories`,索引残留消除(J3,PoC 实证);config/prompt/对话写路径统一原子写(`common/atomic.py`,mkstemp + fsync + os.replace)+ 损坏备份 `config.json.corrupt.<时间戳>` 拒绝覆盖(J5);PromptManager/ConfigManager 单例初始化纳入锁(J6,并发首例不再双重加载系统提示词)
|
||||
- **K4(中危)**:`model_config.py` 连通性验证异常消息复用 `_sanitize_log_msg` 脱敏(J4,Bearer 令牌不再进终端);`broadcast` 改每客户端并发发送 + `asyncio.wait_for` 超时(`broadcast_send_timeout` 默认 5s),超时即 unregister——僵死客户端(半开 TCP)不再卡死广播与调度器循环(J7)
|
||||
- **K5(低危随批)**:J8 交互输入容错(非法温度/超时友好报错);J9 MemoDatabase WAL + busy_timeout=10000 覆盖全部写路径;J10 memo 批量删除补实现(`memo_ids`);J11 `/new` 取消-重置竞态(排队 reset 等待 is_processing 清零);J12 skills 注入面(nmap hosts `shlex.quote`、screenshot PowerShell 单引号转义、SKILLS_FORMAT.md 新增"安全边界"节);J14 voice_client 死代码/僵尸连接/录音残留清理;J15 死代码移除(`full_response_queue` 等)+ 未知 task action 明确报错 + 广播门控对齐 + 热路径日志降级;J16 `--help` 输出 2000 上限 + 同名技能键确定性解决
|
||||
- ✅ **测试**:490 项全过(Linux 环境;第五轮新增 Batch K 35 项 + 子代理 max_turns 回归 14 项;`test_execute_timeout`/`test_flat_api_still_reads_key_file` 为 Windows 向用例,Linux 预存失败与本轮无关)
|
||||
|
||||
### v3.5
|
||||
- ✅ **安全审计与修复(第三轮 Batch G + 第四轮 Batch H + 遗留收尾 Batch I)**(完整记录见 [docs/reports/aibox_audit_20260817.md](../docs/reports/aibox_audit_20260817.md)):
|
||||
- **Batch G(llm 加固 + 线程生命周期 + 取消传播 + DB 事务)**:
|
||||
- llm_client:非瞬时 HTTP 错误体截断 `[:200]`(密钥/错误页不进异常栈)、Retry-After 上限 cap 3600s、fallback 档案缺 model 跳过并告警、fallback 去重、删除 `LLMClient._cancelled` 死状态、FailoverClient 并发契约文档化
|
||||
- 子代理取消传播:`run_full_subagent` 共享父 `_cancel_event` + 0.2s 观察者线程转发深度中断(强关在途 HTTP、kill 子进程);`SubAgentTool` 构造注入父事件
|
||||
- drain 线程重构:批次级单线程 + stop 事件 + try/finally 统一回收(消除每次工具调用泄漏一个常驻线程);`/new` 的 `reset_conversation` 移入线程池(不再同步获取 `_msg_lock` 冻结事件循环);broadcast 收窄 except + unregister 改 discard;`_processing_tasks` 改按连接的任务集合登记(关闭时取消全部在途任务)
|
||||
- DB:FTS5 关键词双引号翻倍转义(`"` → `""`)、两处 `LIMIT` f-string 参数化、`mark_as_triggered` 读写收进单连接 `BEGIN IMMEDIATE` 事务(消灭并发触发竞态);历史索引裸 except 收窄 + 归档文件名同秒加序号
|
||||
- **Batch H(第四轮审计,3 路并行审读 + Linux PoC 6/6 实证)**:
|
||||
- **H1(高危)**:修复 N5 回归——子代理共享父取消事件后首轮 think 无条件 `clear()` 会吞掉父级已置位的取消信号(用户中断 / --once 超时被绕过),共享事件时跳过 clear + 父已取消直接早退
|
||||
- `/reload_prompts` 移入线程池(N3 同类遗漏);技能扫描(串行 `--help` 子进程)启动前预热,不再在首条消息/首个任务时冻结事件循环
|
||||
- 调度器任务级超时(`scheduler.task_timeout`,默认 600s)——AI think 卡死不再停摆整个调度循环
|
||||
- `get_recent_memories(limit=None)` 默认上限(不再全表加载);`memo delete` 改物理删除(此前只标记完成);`PromptManager.reload` 本地构建后原子替换(并发读者不再读到空/半截系统提示词)
|
||||
- read_file 的 sed 参数强转 int(防引号注入);`set_proxy` 落盘失败如实报告;SSE 非 UTF-8 行容错;persona 路径 realpath 校验(软链越界拒绝);1MB 截断无半个字符;compressor 除零防护;InstallSkillTool 指引命令路径引用;ask_user 缺参明确报错
|
||||
- **POSIX 进程组终止**:executor 的 `Popen` 加 `start_new_session` + 组信号 kill(组 ID 在 Popen 后立即捕获)——`(sleep 300 &)` 形态的孙进程持管道不再悬挂 reader 线程/泄漏 fd(实证修复前孙进程残留)
|
||||
- web 工具 scheme 白名单(http/https,`file://` 等拒绝);registry 日志打码改包含匹配(apiKey/passwd 变体);FailoverClient 日志 Bearer 令牌正则脱敏(网关回显 Authorization 头不再进日志)
|
||||
- **Batch I(遗留收尾)**:
|
||||
- **N9**:prompt 写入内容长度上限(`MAX_CONTENT_LENGTH=10000`,identity/soul/status 超长拒绝)——约束持久化注入面
|
||||
- **F5**:情绪子代理取消传播——`run_light_subagent` 改走 stream_chat + 观察者线程,用户中断可即时终止(此前阻塞至请求超时)
|
||||
- 提醒时间格式异常任务按"到期"返回(此前被 SQL 层字符串比较静默丢弃,任务永不执行);`mark_as_triggered` 加 `busy_timeout=10000`(多进程同库写竞争等待而非立即报 locked)
|
||||
- 对话文件原子写(tmp + os.replace,崩溃不留半截 JSON);`load_conversation` realpath 越界防护;WS 命令专用线程池(`/new`、`/reload_prompts` 不再占用默认池);--once 硬超时兜底(软超时后宽限一个周期强杀进程)
|
||||
- PromptManager 单例加锁;`limit=0` 语义修正(None 才不加 LIMIT)
|
||||
- ✅ **测试**:424 项全过(Linux 环境;Batch G 20 + Batch H 26 + Batch I 16 项新增;`test_execute_timeout`/`test_flat_api_still_reads_key_file` 为 Windows 向用例,Linux 预存失败与本轮无关)
|
||||
|
||||
### v3.4
|
||||
- ✅ **多模型档案加固(H1/H2/H3/M1/M2/M3 六项修复)**:
|
||||
- **H1(凭据隔离)**:档案端点环境变量取不到时直接报错,**绝不回退读 `~/.cert/deepseek.key`**(该文件为 DeepSeek 专属凭据,跨厂商复用属于泄漏);key 文件回退仅对无档案的旧 `api` 扁平节生效
|
||||
@@ -917,7 +1008,7 @@ elif command == "/mycommand":
|
||||
- ✅ **测试**:272 项全过(Batch E PoC 回归 + F1 配置脱敏 3 项 + F2 超时清理/线程收尾 5 项新增)
|
||||
|
||||
### v3.2
|
||||
- ✅ **CLI 一次性对话模式(`--once`)**(分三批落地,开发文档见 [docs/reports/aibox_cli_once_mode_dev_doc.md](../docs/reports/aibox_cli_once_mode_dev_doc.md)):
|
||||
- ✅ **CLI 一次性对话模式(`--once`)**(分三批落地):
|
||||
- **Batch A**:`--once QUERY` / `--quiet` / `--once-timeout`;不启动 WebSocket,执行后退出;回答走 stdout、过程走 stderr,可管道组合;SIGINT 与超时复用三层取消机制
|
||||
- **Batch B**:`--no-save-memory`(`enable_memory` 开关干净禁用 MemoryTool)、`--show-reasoning`(思考流默认静默);压缩提示改走 output_callback 不再污染 stdout;退出码 1(工具执行失败)
|
||||
- **Batch C**:SQLite WAL(多 CLI 进程并发写锁缓解)、`--worker-only`(纯执行面 + 精简 system prompt 防人格污染)、`--model` / `--temperature` 覆盖、子代理工具失败标记透传(`_SUBAGENT_TOOL_FAILED_MARKER` → orchestrator stats)
|
||||
|
||||
@@ -227,3 +227,14 @@ except ImportError:
|
||||
```
|
||||
|
||||
**`--help` 要完备**。包含参数说明和使用示例两部分。AI 通过阅读 `--help` 输出学会如何调用技能,信息越完整,AI 调用越准确。
|
||||
|
||||
---
|
||||
|
||||
## 安全边界(J12,审计 2026-08-19)
|
||||
|
||||
**技能即可信代码**。可执行技能由 AI(可被网页内容/用户消息诱导)以子进程方式运行,参数为不可信输入:
|
||||
|
||||
- **参数必须经 shell 引号转义**:凡技能内部把参数拼进 shell 命令/脚本串(subprocess shell、PowerShell、python-nmap 等内部拼串库),一律先 `shlex.quote`(POSIX)或对应语言转义(PowerShell 单引号翻倍)
|
||||
- **禁止把凭据嵌入命令行参数**:命令行参数对同机进程可见(`ps`);密码/令牌应经环境变量或配置文件传递
|
||||
- **ssh/sftp 类技能**:`AutoAddPolicy` 信任任意主机密钥存在 MITM 风险,仅限可信网络使用;生产环境改用 `known_hosts` 校验
|
||||
- 内置技能示例(`skills/`)已按此规范处理;新增技能请遵守同一条例
|
||||
+454
-99
@@ -26,6 +26,7 @@ import locale
|
||||
import asyncio
|
||||
import websockets
|
||||
from abc import ABC, abstractmethod
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import datetime, timedelta
|
||||
from typing import List, Dict, Optional, Any, Tuple, Callable, Union, Set
|
||||
import argparse
|
||||
@@ -54,6 +55,8 @@ from tools import (
|
||||
UpdateIdentityTool, UpdateSoulTool, SaveMemoryAndEndConversationTool,
|
||||
SetProxyTool, GetProxyTool,
|
||||
AskUserTool,
|
||||
# 后台进程管理工具(bg_run/bg_write/bg_list/bg_output/bg_kill,worker 执行面)
|
||||
BgRunTool, BgWriteTool, BgListTool, BgOutputTool, BgKillTool,
|
||||
)
|
||||
|
||||
# 子系统(Wave 1 + Wave 2 抽离):命令执行器 / 提示词管理器 / 数据库 / 技能管理器 / 对话历史 / 对话压缩
|
||||
@@ -118,7 +121,12 @@ class TaskScheduler:
|
||||
prompt_manager=self.ai_manager.prompt_manager if self.ai_manager else None,
|
||||
skill_manager=self._get_shared_skill_manager(),
|
||||
# H3:CLI --model 为进程级覆盖,定时任务实例同样生效
|
||||
model=getattr(self.ai_manager, "model", None) if self.ai_manager else None
|
||||
model=getattr(self.ai_manager, "model", None) if self.ai_manager else None,
|
||||
# T22:精简 profile——无 update_identity/update_soul/set_proxy,
|
||||
# 被污染的任务内容无法在无人值守时段改写人格/代理
|
||||
tool_profile=TOOL_PROFILE_TASK,
|
||||
# T22:后台进程属主隔离——任务 AI 与会话 AI 的 bg 进程互不可杀
|
||||
owner_id=f"task-{id(self):x}"
|
||||
)
|
||||
|
||||
def set_websocket_handler(self, handler):
|
||||
@@ -177,13 +185,20 @@ class TaskScheduler:
|
||||
for memo in pending:
|
||||
try:
|
||||
# 广播任务开始执行
|
||||
if self.ws_handler:
|
||||
if self.ws_handler and self.broadcast_enabled:
|
||||
await self.ws_handler.broadcast_task_status(memo, "executing", i18n.get('task_executing', memo['title'], memo['id']))
|
||||
|
||||
# 【关键修复】执行前最后一次验证
|
||||
# J1(审计 2026-08-19):格式异常按到期执行(Q5 语义),
|
||||
# 不再抛 ValueError 致任务既不触发也不完成、每秒重试热循环
|
||||
if memo.get('reminder_time'):
|
||||
reminder_time = datetime.strptime(memo['reminder_time'], constants.DATETIME_FORMAT)
|
||||
if reminder_time > now:
|
||||
try:
|
||||
reminder_time = datetime.strptime(memo['reminder_time'], constants.DATETIME_FORMAT)
|
||||
except (ValueError, TypeError) as e:
|
||||
self.logger.warning(
|
||||
f"任务 #{memo['id']} 提醒时间格式异常,按到期执行: {e}")
|
||||
reminder_time = None
|
||||
if reminder_time is not None and reminder_time > now:
|
||||
self.logger.warning(f"任务 #{memo['id']} 的提醒时间 ({memo['reminder_time']}) 大于当前时间 ({now_str}),跳过执行")
|
||||
continue
|
||||
|
||||
@@ -195,7 +210,7 @@ class TaskScheduler:
|
||||
result = await self.execute_task_as_conversation(memo, memo_db)
|
||||
|
||||
# 广播任务完成
|
||||
if self.ws_handler:
|
||||
if self.ws_handler and self.broadcast_enabled:
|
||||
await self.ws_handler.broadcast_task_status(memo, "complete", result)
|
||||
|
||||
executed_tasks.append({
|
||||
@@ -210,7 +225,7 @@ class TaskScheduler:
|
||||
self.logger.error(i18n.get('task_failed', memo['title'], str(e)))
|
||||
|
||||
# 广播任务失败
|
||||
if self.ws_handler:
|
||||
if self.ws_handler and self.broadcast_enabled:
|
||||
await self.ws_handler.broadcast_task_status(memo, "failed", str(e))
|
||||
|
||||
executed_tasks.append({
|
||||
@@ -284,17 +299,11 @@ class TaskScheduler:
|
||||
# 获取当前事件循环
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
# 创建一个队列来收集完整的响应
|
||||
full_response_queue = queue.Queue()
|
||||
|
||||
# 定义输出回调函数 - 直接广播,同时收集
|
||||
# 定义输出回调函数 - 直接广播
|
||||
def output_callback(msg_type, content):
|
||||
"""AI输出的回调函数 - 直接广播,同时收集"""
|
||||
# 收集完整响应
|
||||
if msg_type == "complete":
|
||||
full_response_queue.put(content)
|
||||
|
||||
# 直接广播给所有客户端 - 必须带上 task_id 和 execution_id
|
||||
"""AI输出的回调函数 - 直接广播"""
|
||||
# J15(审计 2026-08-19):full_response_queue 收集后从未被消费,
|
||||
# 属死代码,已移除;广播路径保留
|
||||
if self.ws_handler and self.broadcast_enabled:
|
||||
# 使用 run_coroutine_threadsafe 在正确的循环中运行异步函数
|
||||
asyncio.run_coroutine_threadsafe(
|
||||
@@ -314,28 +323,73 @@ class TaskScheduler:
|
||||
# 创建一个 Future 来等待 AI 完成
|
||||
ai_future = loop.create_future()
|
||||
|
||||
def _set_result_if_pending(response):
|
||||
# T10:wait_for 硬超时会 cancel future,之后 run_ai 线程迟到的
|
||||
# set_result/set_exception 会抛 InvalidStateError 打进事件循环
|
||||
if not ai_future.done():
|
||||
ai_future.set_result(response)
|
||||
|
||||
def _set_exception_if_pending(exc):
|
||||
if not ai_future.done():
|
||||
ai_future.set_exception(exc)
|
||||
|
||||
def run_ai():
|
||||
try:
|
||||
# 在单独的线程中运行 AI(独立实例,天然隔离,无并发冲突)
|
||||
response = ai.think_and_respond(user_input, output_callback)
|
||||
# 将结果设置到 Future 中
|
||||
loop.call_soon_threadsafe(ai_future.set_result, response)
|
||||
loop.call_soon_threadsafe(_set_result_if_pending, response)
|
||||
except Exception as e:
|
||||
loop.call_soon_threadsafe(ai_future.set_exception, e)
|
||||
loop.call_soon_threadsafe(_set_exception_if_pending, e)
|
||||
|
||||
# 启动 AI 线程
|
||||
ai_thread = threading.Thread(target=run_ai)
|
||||
ai_thread.daemon = True
|
||||
ai_thread.start()
|
||||
|
||||
# 等待 AI 完成,同时处理可能的取消
|
||||
# M3(审计 2026-08-18):任务级超时兜底。此前 await ai_future
|
||||
# 无任何超时/取消路径——AI think 卡死(LLM 挂起、子进程不可杀)
|
||||
# 时该任务永不返回,调度器 run 循环顺序 await 被阻塞,后续全部
|
||||
# 定时/即时任务不再执行且任务已 mark_as_triggered 不会重试。
|
||||
# 到点 call_later 触发 ai.cancel()(置取消事件 + 关在途 HTTP 流),
|
||||
# worker 线程自然收尾,调度循环继续。
|
||||
task_timeout = float(ConfigManager().get("scheduler.task_timeout", 600))
|
||||
timed_out = threading.Event()
|
||||
|
||||
def _timeout_cancel():
|
||||
timed_out.set()
|
||||
try:
|
||||
ai.cancel()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
timeout_handle = loop.call_later(task_timeout, _timeout_cancel)
|
||||
try:
|
||||
response = await ai_future
|
||||
# T10(审计 2026-08-21):wait_for 硬超时兜底——软取消(ai.cancel
|
||||
# 三层:取消事件/关在途 HTTP 流/杀子进程)够不到的卡死路径
|
||||
# (Windows 代理连接黑洞、requests 读超时被服务器滴流输出
|
||||
# 规避、kill 失败的进程)会让 run_ai 永不回调,await 永久挂起
|
||||
# 停摆整个调度循环且任务已 mark_as_triggered 不会重试。
|
||||
# task_timeout 到点先软取消,+30s 宽限后硬超时放行调度循环
|
||||
# (对齐 --once 的 _hard_exit 兜底语义;AI 线程为 daemon,
|
||||
# 迟到的回调经 done() 检查安全跳过)
|
||||
response = await asyncio.wait_for(
|
||||
ai_future, timeout=task_timeout + 30)
|
||||
result = response or "AI无响应"
|
||||
except asyncio.TimeoutError:
|
||||
self.logger.error(
|
||||
f"任务执行硬超时({task_timeout + 30:.0f}秒,软取消未生效),"
|
||||
f"放弃等待,调度循环继续")
|
||||
result = f"任务执行硬超时({task_timeout + 30:.0f}秒),已放弃等待"
|
||||
except Exception as e:
|
||||
self.logger.error(f"AI对话失败: {e}")
|
||||
result = f"任务执行但AI对话失败: {str(e)}"
|
||||
if timed_out.is_set():
|
||||
self.logger.error(f"任务执行超时({task_timeout}秒),AI 已取消")
|
||||
result = f"任务执行超时({task_timeout}秒),AI 已取消"
|
||||
else:
|
||||
self.logger.error(f"AI对话失败: {e}")
|
||||
result = f"任务执行但AI对话失败: {str(e)}"
|
||||
finally:
|
||||
timeout_handle.cancel()
|
||||
# 确保线程结束
|
||||
ai_thread.join(timeout=1.0)
|
||||
|
||||
@@ -398,7 +452,7 @@ class TaskScheduler:
|
||||
executed_tasks = []
|
||||
for memo in immediate:
|
||||
try:
|
||||
if self.ws_handler:
|
||||
if self.ws_handler and self.broadcast_enabled:
|
||||
await self.ws_handler.broadcast_task_status(memo, "executing", i18n.get('immediate_memo_executing', memo['title']))
|
||||
|
||||
self.logger.info(f"标记即时任务 #{memo['id']} 为已触发")
|
||||
@@ -406,7 +460,7 @@ class TaskScheduler:
|
||||
|
||||
result = await self.execute_task_as_conversation(memo, memo_db)
|
||||
|
||||
if self.ws_handler:
|
||||
if self.ws_handler and self.broadcast_enabled:
|
||||
await self.ws_handler.broadcast_task_status(memo, "complete", result)
|
||||
|
||||
executed_tasks.append({
|
||||
@@ -420,7 +474,7 @@ class TaskScheduler:
|
||||
except Exception as e:
|
||||
self.logger.error(i18n.get('task_failed', memo['title'], str(e)))
|
||||
|
||||
if self.ws_handler:
|
||||
if self.ws_handler and self.broadcast_enabled:
|
||||
await self.ws_handler.broadcast_task_status(memo, "failed", str(e))
|
||||
|
||||
executed_tasks.append({
|
||||
@@ -490,6 +544,12 @@ class TaskScheduler:
|
||||
|
||||
|
||||
class ExpressEmotionTool(Tool):
|
||||
def __init__(self, logger=None, cancel_event: threading.Event = None):
|
||||
super().__init__(logger)
|
||||
# I 批(审计 2026-08-18):父代理取消事件注入——情绪子代理此前
|
||||
# 用独立 client 且不可中断,父级 cancel 后仍阻塞至请求自然结束
|
||||
self.cancel_event = cancel_event
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "express_emotion"
|
||||
|
||||
@@ -534,10 +594,11 @@ class ExpressEmotionTool(Tool):
|
||||
if not name:
|
||||
return ""
|
||||
path = os.path.join(constants.SKILLS_DIR, name + ".md")
|
||||
# N12:persona 为 LLM 可控输入,做路径边界校验——normcase 统一大小写与
|
||||
# 分隔符后要求解析结果仍在 SKILLS_DIR 内,防止 ../../ 越界读取并注入
|
||||
# N12 + L4(审计 2026-08-18):persona 为 LLM 可控输入。abspath 前缀
|
||||
# 检查拦截 ../ 越界,但符号链接可绕过(skills 目录内软链指向任意文件
|
||||
# 即被读取注入);realpath 解析软链后再做前缀校验。
|
||||
skills_root = os.path.normcase(os.path.abspath(constants.SKILLS_DIR))
|
||||
resolved = os.path.normcase(os.path.abspath(path))
|
||||
resolved = os.path.normcase(os.path.realpath(path))
|
||||
if not resolved.startswith(skills_root + os.sep):
|
||||
self.logger.warning(f"人格技能名非法(越界路径): {name}")
|
||||
return ""
|
||||
@@ -572,7 +633,8 @@ class ExpressEmotionTool(Tool):
|
||||
task = "\n".join(task_lines)
|
||||
|
||||
# 情绪子代理:独立上下文/人格,生成情绪化回应
|
||||
return run_light_subagent(role, task, self.config, logger=self.logger, max_tokens=2048)
|
||||
return run_light_subagent(role, task, self.config, logger=self.logger,
|
||||
max_tokens=2048, cancel_event=self.cancel_event)
|
||||
except Exception as e:
|
||||
self.logger.error(f"情绪子代理执行异常: {e}")
|
||||
return f"❌ 情绪子代理执行失败: {str(e)}"
|
||||
@@ -584,12 +646,20 @@ class ExpressEmotionTool(Tool):
|
||||
|
||||
|
||||
def run_light_subagent(role, task, config, logger=None, model=None,
|
||||
temperature=None, max_tokens=2048) -> str:
|
||||
temperature=None, max_tokens=2048,
|
||||
cancel_event: threading.Event = None) -> str:
|
||||
"""轻量子代理执行:独立上下文/角色,无工具,一次 LLM 补全返回结果。
|
||||
|
||||
供 express_emotion 情绪子代理使用。
|
||||
返回子代理生成的完整文本;失败时返回带 ❌ 的错误信息。
|
||||
|
||||
cancel_event(I 批,审计 2026-08-18):父代理取消事件。此前本函数用
|
||||
非流式 chat_completion,父级 cancel 只能等请求自然结束(最长 timeout,
|
||||
期间持锁阻塞);改走 stream_chat 后由观察者线程转发 sub_client.cancel()
|
||||
强关在途 HTTP 流,用户中断可即时终止。
|
||||
"""
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
return "❌ 子代理已因父级取消而终止"
|
||||
try:
|
||||
# 幂等校验:限制 max_tokens 范围,防止意外超限大量消耗
|
||||
max_tokens = max(1, min(int(max_tokens or 2048), constants.DEFAULT_MAX_TOKENS))
|
||||
@@ -605,8 +675,33 @@ def run_light_subagent(role, task, config, logger=None, model=None,
|
||||
]
|
||||
if logger is not None:
|
||||
logger.info(f"子代理执行: model={model or sub_client.model}, role={str(role)[:40]}...")
|
||||
result = sub_client.chat_completion(messages, temperature=temp, max_tokens=max_tokens)
|
||||
if result is None or not str(result).strip():
|
||||
|
||||
_watch_done = threading.Event()
|
||||
if cancel_event is not None:
|
||||
def _watch_parent_cancel():
|
||||
while not _watch_done.wait(timeout=0.2):
|
||||
if cancel_event.is_set():
|
||||
try:
|
||||
sub_client.cancel()
|
||||
except Exception:
|
||||
pass
|
||||
return
|
||||
|
||||
threading.Thread(target=_watch_parent_cancel, daemon=True,
|
||||
name="light-subagent-cancel-watch").start()
|
||||
|
||||
# I 批:改走流式(可被 cancel 强关在途 HTTP),拼接 content 事件
|
||||
parts = []
|
||||
try:
|
||||
for event in sub_client.stream_chat(messages, temperature=temp,
|
||||
max_tokens=max_tokens):
|
||||
if event.get("type") == "content":
|
||||
parts.append(event.get("text") or "")
|
||||
finally:
|
||||
_watch_done.set()
|
||||
|
||||
result = "".join(parts)
|
||||
if not str(result).strip():
|
||||
return "❌ 子代理执行失败:未获得有效响应"
|
||||
return str(result).strip()
|
||||
except Exception as e:
|
||||
@@ -640,6 +735,9 @@ def run_full_subagent(role, task, config, logger=None, model=None,
|
||||
_watch_done = threading.Event() # N5:取消观察者线程退出信号(finally 置位)
|
||||
try:
|
||||
max_tokens = max(1, min(int(max_tokens or 2048), constants.DEFAULT_MAX_TOKENS))
|
||||
# max_turns 钳制:0/负数此前致首轮即拦截(只剩收尾报告、无任何执行);
|
||||
# 0/空回退默认值的惯例对齐 max_tokens,下限 1 保证至少执行一个任务轮
|
||||
max_turns = max(1, int(max_turns or 10))
|
||||
temp = temperature if temperature is not None else constants.DEFAULT_TEMPERATURE
|
||||
|
||||
# 创建子代理:完整 DeepSeekChat 实例,role 作为 system_prompt
|
||||
@@ -673,6 +771,14 @@ def run_full_subagent(role, task, config, logger=None, model=None,
|
||||
# 深度中断(sub_chat.cancel() 强关在途 HTTP 流并 kill 其子进程)。
|
||||
if cancel_event is not None:
|
||||
sub_chat._cancel_event = cancel_event
|
||||
# H1(审计 2026-08-18):共享父事件后禁止子代理首轮 clear——
|
||||
# _think_and_respond_impl 首轮无条件 clear() 会吞掉父级已置位的
|
||||
# 取消信号(用户中断 / --once 超时恰落在子代理生成窗口时被清掉,
|
||||
# 观察者线程随之失效,子代理跑完全程)。置位后立即检查一次,
|
||||
# 父级已取消则不再启动子代理。
|
||||
sub_chat._shared_cancel_event = True
|
||||
if cancel_event.is_set():
|
||||
return "❌ 子代理已因父级取消而终止"
|
||||
|
||||
def _watch_parent_cancel():
|
||||
while not _watch_done.wait(timeout=0.2):
|
||||
@@ -910,9 +1016,14 @@ class SubAgentTool(Tool):
|
||||
|
||||
TOOL_PROFILE_ORCHESTRATOR = "orchestrator"
|
||||
TOOL_PROFILE_WORKER = "worker"
|
||||
# T22(审计 2026-08-21):定时任务精简 profile——任务内容可源自 LLM 生成物
|
||||
# (备忘录标题/正文可能被网页内容污染),无人值守时段不持有改写全局人格
|
||||
# 提示词(update_identity/update_soul)与全局网络面(set_proxy)的工具
|
||||
TOOL_PROFILE_TASK = "task"
|
||||
|
||||
# 仅 worker 子代理注册的执行面工具;编排者误调时拦截并纠正
|
||||
WORKER_ONLY_TOOLS = {"run_command", "read_file", "write_file", "list_files", "web", "calculator"}
|
||||
WORKER_ONLY_TOOLS = {"run_command", "read_file", "write_file", "list_files", "web", "calculator",
|
||||
"bg_run", "bg_write", "bg_list", "bg_output", "bg_kill"}
|
||||
|
||||
# 每轮临时注入的互动风格提醒(仅编排者,不进对话历史)
|
||||
ORCHESTRATOR_TURN_GUIDANCE = (
|
||||
@@ -947,12 +1058,16 @@ class DeepSeekChat:
|
||||
memo_db: MemoDatabase = None, command_executor: CommandExecutor = None,
|
||||
prompt_manager: PromptManager = None, skill_manager: SkillManager = None,
|
||||
tool_profile: str = TOOL_PROFILE_ORCHESTRATOR, enable_memory: bool = True,
|
||||
model: str = None, temperature: float = None):
|
||||
model: str = None, temperature: float = None,
|
||||
owner_id: str = None):
|
||||
self.config = ConfigManager()
|
||||
self.logger = Logger("DeepSeek", debug_level if debug_level is not None else Logger("").level)
|
||||
self.debug_level = self.logger.level
|
||||
# 工具集 profile:决定本实例注册哪套工具(orchestrator 管理面 / worker 执行面)
|
||||
self.tool_profile = tool_profile
|
||||
# T22:后台进程属主标识——bg_* 工具携带,ProcessManager 据此隔离
|
||||
# kill/write 权限(会话/任务/子代理互不可操作他人进程,观察面共享)
|
||||
self.owner_id = owner_id
|
||||
# 大模型客户端(抽象层,支持多厂商切换);model 显式传入时覆盖配置默认
|
||||
self.client = create_llm_client(self.config, api_key=api_key, logger=self.logger, model=model)
|
||||
self.api_key = self.client.api_key
|
||||
@@ -1104,12 +1219,22 @@ class DeepSeekChat:
|
||||
SimpleWebCrawlerTool(self.logger),
|
||||
CalculatorTool(self.logger),
|
||||
GetCurrentTimeTool(self.logger),
|
||||
# 后台进程管理(注册表全局共享;T22:owner 隔离写/杀权限,
|
||||
# 观察面 list/output 仍全局可见)
|
||||
BgRunTool(self.logger, owner=self.owner_id),
|
||||
BgWriteTool(self.logger, owner=self.owner_id),
|
||||
BgListTool(self.logger),
|
||||
BgOutputTool(self.logger),
|
||||
BgKillTool(self.logger, owner=self.owner_id),
|
||||
]
|
||||
self.tool_registry.register_many(worker_tools, is_builtin=True)
|
||||
self.logger.info(f"worker 子代理已注册 {len(worker_tools)} 个执行面工具")
|
||||
return
|
||||
|
||||
# orchestrator(主对话):管理面工具,执行类任务经 create_subagent 派生 worker
|
||||
# T22:task profile 为 orchestrator 子集——剔除改写全局人格提示词与
|
||||
# 全局网络面的工具(定时任务无人值守,任务内容可能被污染)
|
||||
is_task_profile = self.tool_profile == TOOL_PROFILE_TASK
|
||||
builtin_tools = [
|
||||
GetCurrentTimeTool(self.logger),
|
||||
]
|
||||
@@ -1120,26 +1245,30 @@ class DeepSeekChat:
|
||||
if self.memo_db:
|
||||
builtin_tools.append(MemoTool(self.memo_db, self.logger))
|
||||
|
||||
# 更新身份工具
|
||||
builtin_tools.append(UpdateIdentityTool(self.prompt_manager, self.logger))
|
||||
if not is_task_profile:
|
||||
# 更新身份工具
|
||||
builtin_tools.append(UpdateIdentityTool(self.prompt_manager, self.logger))
|
||||
|
||||
# 新增:更新智能体自我认识工具
|
||||
builtin_tools.append(UpdateSoulTool(self.prompt_manager, self.logger))
|
||||
# 新增:更新智能体自我认识工具
|
||||
builtin_tools.append(UpdateSoulTool(self.prompt_manager, self.logger))
|
||||
|
||||
# 增强版保存记忆并结束对话工具
|
||||
if exit_callback and self.memory_db:
|
||||
builtin_tools.append(SaveMemoryAndEndConversationTool(exit_callback, self.memory_db, self.prompt_manager, self.logger))
|
||||
|
||||
# 添加询问用户的工具
|
||||
# I 批:ExpressEmotionTool 注入父取消事件(情绪子代理可被中断)
|
||||
builtin_tools.append(AskUserTool(self.logger))
|
||||
builtin_tools.append(ExpressEmotionTool(self.logger))
|
||||
builtin_tools.append(ExpressEmotionTool(self.logger, cancel_event=self._cancel_event))
|
||||
|
||||
# 添加技能热重载工具 & 技能安装指引工具 & 创建技能工具
|
||||
builtin_tools.append(ReloadSkillsTool(self.tool_registry, self.logger))
|
||||
builtin_tools.append(InstallSkillTool(self.logger))
|
||||
builtin_tools.append(CreateSkillTool(self.logger))
|
||||
|
||||
builtin_tools.append(SetProxyTool(self.logger))
|
||||
if not is_task_profile:
|
||||
# T22:set_proxy 是全局持久化网络面,任务 AI 不持有
|
||||
builtin_tools.append(SetProxyTool(self.logger))
|
||||
builtin_tools.append(GetProxyTool(self.logger))
|
||||
|
||||
# 子代理工具:主代理可派生独立上下文/角色的轻量子代理执行任务
|
||||
@@ -1291,13 +1420,20 @@ class DeepSeekChat:
|
||||
# 队列是否为空),闭包长期持有 output_callback(含 websocket 引用)
|
||||
# 阻碍 GC,常驻服务下线程与内存持续累积。
|
||||
drain_stop = threading.Event()
|
||||
# T12:drain 线程与工具调用后的排空循环并发消费同一队列,complete
|
||||
# 标记可能被任一方取走——用可变容器共享,两路见到都置位,供
|
||||
# has_stream_output 判定流式工具(流式已逐行转发,勿重复整发结果)
|
||||
saw_complete = [False]
|
||||
|
||||
def _drain_queue(q=output_queue, cb=output_callback, stop_event=drain_stop):
|
||||
def _drain_queue(q=output_queue, cb=output_callback, stop_event=drain_stop,
|
||||
flag=saw_complete):
|
||||
"""持续排空 output_queue 并转发 line 输出(长时间工具执行期间防 WS 空闲断线)"""
|
||||
while not stop_event.is_set():
|
||||
try:
|
||||
msg_type, content = q.get(timeout=0.1)
|
||||
if cb and msg_type == "line":
|
||||
if msg_type == "complete":
|
||||
flag[0] = True
|
||||
elif cb and msg_type == "line":
|
||||
cb("line", str(content) if content is not None else "")
|
||||
except queue.Empty:
|
||||
continue
|
||||
@@ -1329,8 +1465,8 @@ class DeepSeekChat:
|
||||
output_callback("line", f"🔧 调用工具: {tool_name}\n")
|
||||
|
||||
result_raw = None
|
||||
if self.tool_profile == TOOL_PROFILE_ORCHESTRATOR and tool_name in WORKER_ONLY_TOOLS:
|
||||
# 编排者不持有执行面工具:拦截并纠正,引导改用 create_subagent
|
||||
if self.tool_profile in (TOOL_PROFILE_ORCHESTRATOR, TOOL_PROFILE_TASK) and tool_name in WORKER_ONLY_TOOLS:
|
||||
# 编排者/任务 AI 不持有执行面工具:拦截并纠正,引导改用 create_subagent
|
||||
result_raw = (
|
||||
f"⛔ '{tool_name}' 调用被拦截:你是编排者,没有执行面工具(命令/文件/网页由 worker 子代理持有)。\n"
|
||||
"请改用 create_subagent 派生子代理执行该任务:在 role 中写明执行者身份,"
|
||||
@@ -1338,6 +1474,9 @@ class DeepSeekChat:
|
||||
)
|
||||
self.logger.warning(f"编排者误调执行面工具 {tool_name},已拦截并引导改用 create_subagent")
|
||||
else:
|
||||
# T12:每个工具调用前重置流式标志(drain 线程与下方排空
|
||||
# 循环并发消费,任一方见到 complete 都会置位)
|
||||
saw_complete[0] = False
|
||||
try:
|
||||
result_raw = self.tool_registry.execute_tool(tool_name, **args)
|
||||
# 排空队列中剩余标记(complete 标记供 has_stream_output 判定、
|
||||
@@ -1345,7 +1484,9 @@ class DeepSeekChat:
|
||||
while not output_queue.empty():
|
||||
try:
|
||||
msg_type, content = output_queue.get_nowait()
|
||||
if output_callback and msg_type == "line":
|
||||
if msg_type == "complete":
|
||||
saw_complete[0] = True
|
||||
elif output_callback and msg_type == "line":
|
||||
output_callback("line", str(content) if content is not None else "")
|
||||
except queue.Empty:
|
||||
break
|
||||
@@ -1363,20 +1504,20 @@ class DeepSeekChat:
|
||||
# 安全转换并清理结果
|
||||
result = self._sanitize_tool_result(result_raw)
|
||||
|
||||
# drain 线程持续转发 line 输出,这里通过 complete 标记判断是否流式工具
|
||||
has_stream_output = not output_queue.empty()
|
||||
# T12:流式工具的 line 输出已由 drain 线程逐行转发(registry 侧
|
||||
# N8 预算已生效),勿再把完整 result 整发——此前排空后判定
|
||||
# output_queue.empty() 恒为 False,导致重复推送 + 超长结果绕过预算
|
||||
has_stream_output = saw_complete[0]
|
||||
|
||||
# 普通(非流式)工具没有队列输出,直接发送结果
|
||||
# 长度预算前置:超长结果先替换为提示,再决定是否推送,
|
||||
# 避免全量 result 经非流式直发路径绕过 N8 输出限制(T12)
|
||||
if len(result) > MAX_RESULT_LENGTH:
|
||||
result = f"⚠️ 工具调用结果长度超过 {MAX_RESULT_LENGTH} 字符,已丢弃。请使用更精确的查询获取所需信息。"
|
||||
|
||||
# 普通(非流式)工具没有队列输出,直接发送(已截断的)结果
|
||||
if not has_stream_output and output_callback and result:
|
||||
output_callback("line", result + "\n")
|
||||
|
||||
# 检查结果大小,如果超过阈值,不保留截断内容,仅返回提示信息
|
||||
if len(result) > MAX_RESULT_LENGTH:
|
||||
truncated_note = f"⚠️ 工具调用结果长度超过 {MAX_RESULT_LENGTH} 字符,已丢弃。请使用更精确的查询获取所需信息。"
|
||||
result = truncated_note
|
||||
if output_callback:
|
||||
output_callback("line", truncated_note + "\n")
|
||||
|
||||
# 结束调用工具时发送换行
|
||||
if output_callback:
|
||||
output_callback("line", "\n\n")
|
||||
@@ -1456,8 +1597,9 @@ class DeepSeekChat:
|
||||
|
||||
def _think_and_respond_impl(self, input_text: str, output_callback: Callable = None, is_continuation: bool = False) -> Optional[str]:
|
||||
"""think_and_respond 的实际实现,在 _msg_lock 保护下执行。"""
|
||||
# 首轮开始时清除取消标志
|
||||
if not is_continuation:
|
||||
# 首轮开始时清除取消标志(H1:共享父事件的子代理不清除——清除会
|
||||
# 吞掉父级已置位的取消信号,使中断/--once 超时失效)
|
||||
if not is_continuation and not getattr(self, "_shared_cancel_event", False):
|
||||
self._cancel_event.clear()
|
||||
|
||||
# 立即检查是否已被取消
|
||||
@@ -2021,6 +2163,19 @@ class WebSocketHandler:
|
||||
self.client_sessions = {}
|
||||
self._processing_tasks = {} # websocket -> {asyncio.Task}(N14:集合登记,支持并发多任务)
|
||||
self._shared_skill_manager = None # 共享技能管理器(懒加载,供所有连接复用)
|
||||
# I 批(审计 2026-08-18):/new、/reload_prompts 等命令专用线程池——
|
||||
# 此前 run_in_executor(None) 与 AI think 共用默认池,繁忙时 reset 任务
|
||||
# 排队等锁,多连接并发 /new 可耗尽默认池(阻塞新消息处理)
|
||||
self._cmd_executor = ThreadPoolExecutor(max_workers=2,
|
||||
thread_name_prefix="ws-cmd")
|
||||
# J7(审计 2026-08-19):broadcast 单客户端发送超时(半开 TCP 防卡死)
|
||||
self.broadcast_send_timeout = 5.0
|
||||
|
||||
def shutdown(self):
|
||||
"""I 批:释放命令专用线程池(服务退出时调用)"""
|
||||
if self._cmd_executor is not None:
|
||||
self._cmd_executor.shutdown(wait=False)
|
||||
self._cmd_executor = None
|
||||
|
||||
def _get_shared_skill_manager(self):
|
||||
"""获取共享技能管理器(懒加载一次,供所有连接实例复用,避免重复扫描技能)"""
|
||||
@@ -2082,6 +2237,14 @@ class WebSocketHandler:
|
||||
# N13:discard 替代 remove——并发注销/重复注销时 remove 抛 KeyError
|
||||
# 会逃逸到调用方(如 broadcast 的 except 分支),discard 幂等
|
||||
self.connected_clients.discard(websocket)
|
||||
# R3(审计 2026-08-20):取消该连接仍在进行的处理任务并主动关闭
|
||||
# websocket。broadcast 超时路径(J7)也会走到这里——ping 禁用下
|
||||
# 半开 TCP 的接收循环对断开无感知,不 close 的话连接对象/协程/
|
||||
# 任务登记泄漏直至进程重启;与 handle_connection finally 对齐
|
||||
pending = self._processing_tasks.pop(websocket, set())
|
||||
for t in pending:
|
||||
if not t.done():
|
||||
t.cancel()
|
||||
if websocket in self.client_sessions:
|
||||
session = self.client_sessions.pop(websocket)
|
||||
# 释放该连接独立的AI实例资源
|
||||
@@ -2091,6 +2254,11 @@ class WebSocketHandler:
|
||||
ai.cancel()
|
||||
except Exception:
|
||||
pass
|
||||
# close 幂等(已关闭连接重复调用无副作用);mock/异常对象兜底
|
||||
try:
|
||||
await websocket.close()
|
||||
except Exception:
|
||||
pass
|
||||
self.logger.info(i18n.get('client_disconnected', websocket.remote_address))
|
||||
|
||||
async def send_welcome_info(self, websocket):
|
||||
@@ -2185,14 +2353,31 @@ class WebSocketHandler:
|
||||
self.logger.error(f"发送消息失败: {e}")
|
||||
|
||||
async def broadcast(self, message):
|
||||
"""广播消息到所有客户端"""
|
||||
for websocket in self.connected_clients.copy():
|
||||
"""广播消息到所有客户端。
|
||||
|
||||
J7(审计 2026-08-19):此前串行 await 且 send_message 吞掉全部异常,
|
||||
broadcast 的 except→unregister 为死代码——ping 禁用(ping_interval=None)
|
||||
下半开 TCP 无法探测,backpressure 下 websocket.send 可无限期阻塞,
|
||||
卡死广播与调度器循环。现改为每客户端并发发送 + wait_for 超时,
|
||||
超时即 unregister 回收(discard 幂等)。
|
||||
"""
|
||||
clients = list(self.connected_clients)
|
||||
if not clients:
|
||||
return
|
||||
|
||||
async def _send_one(websocket):
|
||||
try:
|
||||
await self.send_message(websocket, message)
|
||||
await asyncio.wait_for(
|
||||
self.send_message(websocket, message),
|
||||
timeout=self.broadcast_send_timeout)
|
||||
except Exception:
|
||||
# N13:不再裸 except——吞掉 CancelledError 会破坏任务取消语义
|
||||
# 超时/发送异常:回收该连接(取消事件被 wait_for 传播,
|
||||
# CancelledError 为 BaseException 不在此列,不影响任务取消语义)
|
||||
await self.unregister(websocket)
|
||||
|
||||
await asyncio.gather(*(_send_one(ws) for ws in clients),
|
||||
return_exceptions=True)
|
||||
|
||||
async def broadcast_task_status(self, task: Dict, status: str, message: str = None):
|
||||
"""广播任务状态"""
|
||||
await self.broadcast({
|
||||
@@ -2253,21 +2438,24 @@ class WebSocketHandler:
|
||||
"type": "interrupt_ack",
|
||||
"content": "已收到中断信号,正在停止..."
|
||||
})
|
||||
# 取消当前连接独立AI实例的对话(线程安全,只set Event + close连接 + kill进程)
|
||||
# 取消当前连接独立AI实例的对话(线程安全,只set Event + close连接 + kill进程);
|
||||
# ai.cancel() 已同时终止该会话自身 executor 中的子进程。
|
||||
# R2(审计 2026-08-20):移除对 ai_manager.command_executor(CLI 主 AI
|
||||
# 执行器)的全局取消——会话隔离架构下单实例时代的遗留,WS 用户中断
|
||||
# 会误杀终端用户正在执行的命令;对会话自身亦冗余
|
||||
session = self._get_session(websocket)
|
||||
ai = session.get("ai") if session else None
|
||||
if ai:
|
||||
ai.cancel()
|
||||
# 同时确保命令执行器中的进程也被终止
|
||||
if self.ai_manager.command_executor and hasattr(self.ai_manager.command_executor, 'cancel'):
|
||||
try:
|
||||
self.ai_manager.command_executor.cancel()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def handle_task_message(self, websocket, data):
|
||||
"""处理任务相关消息 - 简化版,所有任务都是对话任务"""
|
||||
action = data.get("action", "").strip()
|
||||
# W1(审计 2026-08-20):字段类型规整——非 str 此前在 .strip() 抛
|
||||
# AttributeError,落顶层兜底以笼统"处理失败"收场
|
||||
action = data.get("action", "")
|
||||
if not isinstance(action, str):
|
||||
action = ""
|
||||
action = action.strip()
|
||||
|
||||
if action == "add":
|
||||
# 添加任务
|
||||
@@ -2282,12 +2470,32 @@ class WebSocketHandler:
|
||||
repeat_interval_unit = data.get("repeat_interval_unit", "days")
|
||||
repeat_end_time = data.get("repeat_end_time")
|
||||
|
||||
if not title:
|
||||
# W1(审计 2026-08-20):title 传 dict 等类型此前直达 sqlite 抛
|
||||
# InterfaceError(J1 except ValueError 拦不住)
|
||||
if not isinstance(title, str) or not title.strip():
|
||||
await self.send_message(websocket, {
|
||||
"type": "error",
|
||||
"content": "需要提供任务标题"
|
||||
})
|
||||
return
|
||||
if content is not None and not isinstance(content, str):
|
||||
await self.send_message(websocket, {
|
||||
"type": "error",
|
||||
"content": "任务内容必须是字符串"
|
||||
})
|
||||
return
|
||||
|
||||
# W2(审计 2026-08-20):delay 无校验——非数字 TypeError、负数立即
|
||||
# 执行、10^18 OverflowError。float 转换 + 钳制 [0, 10 年]
|
||||
try:
|
||||
delay = float(delay)
|
||||
except (TypeError, ValueError):
|
||||
await self.send_message(websocket, {
|
||||
"type": "error",
|
||||
"content": f"任务延迟必须是数字(秒): {delay!r}"
|
||||
})
|
||||
return
|
||||
delay = min(max(delay, 0.0), 10 * 365 * 86400)
|
||||
|
||||
reminder_time = None
|
||||
if is_immediate:
|
||||
@@ -2295,16 +2503,26 @@ class WebSocketHandler:
|
||||
else:
|
||||
reminder_time = (datetime.now() + timedelta(seconds=delay)).strftime(constants.DATETIME_FORMAT)
|
||||
|
||||
memo_id = self.ai_manager.memo_db.add_user_task(
|
||||
title=title,
|
||||
content=content,
|
||||
reminder_time=reminder_time,
|
||||
repeat_type=repeat_type,
|
||||
repeat_interval=repeat_interval,
|
||||
repeat_interval_unit=repeat_interval_unit,
|
||||
repeat_end_time=repeat_end_time,
|
||||
is_immediate=is_immediate
|
||||
)
|
||||
# J1(审计 2026-08-19):add_memo 入库校验抛 ValueError(畸形时间/
|
||||
# 重复间隔等),此处转为友好错误消息,不让 handle_message 顶层
|
||||
# 兜底以笼统的"处理失败"收场
|
||||
try:
|
||||
memo_id = self.ai_manager.memo_db.add_user_task(
|
||||
title=title,
|
||||
content=content,
|
||||
reminder_time=reminder_time,
|
||||
repeat_type=repeat_type,
|
||||
repeat_interval=repeat_interval,
|
||||
repeat_interval_unit=repeat_interval_unit,
|
||||
repeat_end_time=repeat_end_time,
|
||||
is_immediate=is_immediate
|
||||
)
|
||||
except ValueError as e:
|
||||
await self.send_message(websocket, {
|
||||
"type": "error",
|
||||
"content": f"任务添加失败: {e}"
|
||||
})
|
||||
return
|
||||
|
||||
immediate_msg = i18n.get('immediate_memo_added') if is_immediate else ""
|
||||
repeat_msg = ""
|
||||
@@ -2353,10 +2571,25 @@ class WebSocketHandler:
|
||||
"type": "error",
|
||||
"content": i18n.get('memo_complete_failed', memo_id)
|
||||
})
|
||||
else:
|
||||
# W3(审计 2026-08-20):缺 task_id 此前静默无响应
|
||||
await self.send_message(websocket, {
|
||||
"type": "error",
|
||||
"content": i18n.get('memo_complete_missing_id')
|
||||
})
|
||||
else:
|
||||
# J15(审计 2026-08-19):未知 action 此前静默吞掉,客户端无任何反馈
|
||||
await self.send_message(websocket, {
|
||||
"type": "error",
|
||||
"content": f"未知的任务操作: {action}"
|
||||
})
|
||||
|
||||
async def handle_command(self, websocket, data):
|
||||
"""处理命令消息"""
|
||||
command = data.get("command", "").strip()
|
||||
command = data.get("command", "")
|
||||
if not isinstance(command, str): # W1(审计 2026-08-20):非 str 字段规整
|
||||
return
|
||||
command = command.strip()
|
||||
|
||||
if not command:
|
||||
return
|
||||
@@ -2383,6 +2616,14 @@ class WebSocketHandler:
|
||||
await asyncio.sleep(0.1)
|
||||
if not ai.is_busy():
|
||||
break
|
||||
# J11(审计 2026-08-19):取消后旧回合可能仍在收尾(回发内容),
|
||||
# 且新消息的 think 线程可能与 reset 无序抢 _msg_lock(RLock 无
|
||||
# 公平性)——等待该连接处理标志清零后再排队 reset,避免新回合
|
||||
# 在旧上下文上执行后被 reset 抹掉(消息"消失")
|
||||
for _ in range(20):
|
||||
if not session or not session.get("is_processing"):
|
||||
break
|
||||
await asyncio.sleep(0.1)
|
||||
# 重置对话状态,重新加载系统提示词
|
||||
# N3:reset_conversation 内部同步获取 _msg_lock(繁忙对话持锁
|
||||
# 可达数十秒),直接调用会冻结整个事件循环(所有连接的心跳/
|
||||
@@ -2394,7 +2635,7 @@ class WebSocketHandler:
|
||||
# 重新初始化提示词管理器(确保最新)
|
||||
self.ai_manager.prompt_manager.reload()
|
||||
|
||||
await loop.run_in_executor(None, _do_reset)
|
||||
await loop.run_in_executor(self._cmd_executor, _do_reset)
|
||||
# 重置该连接对话进行中标记
|
||||
if session:
|
||||
session["dialogue_in_progress"] = False
|
||||
@@ -2441,7 +2682,16 @@ class WebSocketHandler:
|
||||
memos = self.ai_manager.memo_db.get_all_memos(include_completed=False)
|
||||
result = f"\n{i18n.get('pending_memos_title')}\n"
|
||||
for memo in memos:
|
||||
status = i18n.get('memo_overdue') if memo['reminder_time'] and datetime.strptime(memo['reminder_time'], constants.DATETIME_FORMAT) < datetime.now() else i18n.get('memo_pending')
|
||||
# R1(审计 2026-08-20):畸形 reminder_time 此前裸 strptime 抛
|
||||
# ValueError → /memos 整体落入顶层兜底"处理失败"(J1 加固漏点)。
|
||||
# 解析失败按"待执行"展示,对齐 CLI --memo-list 的既有模式
|
||||
status = i18n.get('memo_pending')
|
||||
if memo['reminder_time']:
|
||||
try:
|
||||
if datetime.strptime(memo['reminder_time'], constants.DATETIME_FORMAT) < datetime.now():
|
||||
status = i18n.get('memo_overdue')
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
immediate_tag = "⚡" if memo.get('is_immediate') else ""
|
||||
repeat_info = ""
|
||||
if memo['repeat_type'] != 'none':
|
||||
@@ -2458,6 +2708,13 @@ class WebSocketHandler:
|
||||
|
||||
elif command.startswith("/search "):
|
||||
query = command[8:].strip()
|
||||
# W4(审计 2026-08-20):空关键词此前 LIKE %% 全表兜底刷屏,回用法提示
|
||||
if not query:
|
||||
await self.send_message(websocket, {
|
||||
"type": "command_result",
|
||||
"content": i18n.get('search_usage')
|
||||
})
|
||||
return
|
||||
results = self.ai_manager.memory_db.search_memories([query])
|
||||
result = i18n.get('search_results', query, len(results)) + "\n"
|
||||
for mem in results:
|
||||
@@ -2484,18 +2741,36 @@ class WebSocketHandler:
|
||||
|
||||
elif command == "/reload":
|
||||
self.ai_manager.config.reload()
|
||||
# W5(审计 2026-08-20):reload 经 _load_config 重建 _config,会静默
|
||||
# 丢弃 CLI 进程级覆盖(--debug/--port/--no-scheduler 等)——此处重放
|
||||
# main() 暂存在 config._cli_overrides 的覆盖
|
||||
overrides = getattr(self.ai_manager.config, "_cli_overrides", None)
|
||||
note = ""
|
||||
if overrides:
|
||||
for section, kv in overrides.items():
|
||||
target = self.ai_manager.config._config.get(section)
|
||||
if isinstance(target, dict):
|
||||
target.update(kv)
|
||||
note = i18n.get('config_reloaded_cli_kept')
|
||||
await self.send_message(websocket, {
|
||||
"type": "command_result",
|
||||
"content": i18n.get('config_reloaded')
|
||||
"content": i18n.get('config_reloaded') + note
|
||||
})
|
||||
|
||||
elif command == "/reload_prompts":
|
||||
session = self._get_session(websocket)
|
||||
ai = session.get("ai") if session else None
|
||||
if ai:
|
||||
new_count = ai.reload_prompts()
|
||||
else:
|
||||
new_count = self.ai_manager.prompt_manager.reload()
|
||||
# M1(审计 2026-08-18):reload_prompts 内部同步获取 _msg_lock
|
||||
# (繁忙对话持锁可达数十秒),直接调用会冻结整个事件循环——
|
||||
# N3 只修了 /new,这里同类遗漏。移入线程池执行。
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
def _do_reload_prompts():
|
||||
if ai:
|
||||
return ai.reload_prompts()
|
||||
return self.ai_manager.prompt_manager.reload()
|
||||
|
||||
new_count = await loop.run_in_executor(self._cmd_executor, _do_reload_prompts)
|
||||
await self.send_message(websocket, {
|
||||
"type": "command_result",
|
||||
"content": i18n.get('prompts_reloaded', new_count)
|
||||
@@ -2579,7 +2854,10 @@ class WebSocketHandler:
|
||||
|
||||
async def handle_user_message(self, websocket, data):
|
||||
"""处理用户消息 - 实时流式输出(每个连接使用独立的AI实例,支持多用户并发)"""
|
||||
content = data.get("content", "").strip()
|
||||
content = data.get("content", "")
|
||||
if not isinstance(content, str): # W1(审计 2026-08-20):非 str 字段规整
|
||||
return
|
||||
content = content.strip()
|
||||
|
||||
if not content:
|
||||
return
|
||||
@@ -2625,7 +2903,7 @@ class WebSocketHandler:
|
||||
# 定义输出回调函数 - 实时发送
|
||||
def output_callback(msg_type, content):
|
||||
"""AI输出的回调函数 - 实时发送每个chunk"""
|
||||
self.logger.info(f"输出回调被调用: {msg_type}, 内容长度: {len(content)}")
|
||||
self.logger.debug(f"输出回调被调用: {msg_type}, 内容长度: {len(content)}")
|
||||
|
||||
# 创建发送消息的协程
|
||||
async def send_msg():
|
||||
@@ -2847,8 +3125,11 @@ class EnhancedAIManager:
|
||||
async def handle_connection(self, websocket, path):
|
||||
"""处理WebSocket连接 - 消息接收与处理解耦,确保中断信号能被即时处理"""
|
||||
self.logger.info(f"新连接: {websocket.remote_address}, path: {path}")
|
||||
await self.ws_handler.register(websocket)
|
||||
try:
|
||||
# J1(审计 2026-08-19):register 移入 try——此前在 try 外,握手期
|
||||
# 异常(如任务展示路径解析失败)逃逸时 finally 不执行,会话登记
|
||||
# 永久泄漏;unregister 为 discard 幂等,提前失败也会被正常清理
|
||||
await self.ws_handler.register(websocket)
|
||||
async for message in websocket:
|
||||
# 【关键修复】用create_task异步处理消息,不await阻塞接收循环
|
||||
# 这样当AI在后台生成响应时,interrupt消息仍能被立即接收和处理
|
||||
@@ -2883,12 +3164,17 @@ class EnhancedAIManager:
|
||||
port = constants.WEBSOCKET_PORT
|
||||
|
||||
self.ws_handler = WebSocketHandler(self)
|
||||
# M2(审计 2026-08-18):技能扫描对每个可执行技能串行跑 --help
|
||||
# 子进程(各 ≤10s),懒加载于首条消息/首个任务时在事件循环线程
|
||||
# 执行会冻结整个服务。启动前预热,把扫描成本移到连接建立之前。
|
||||
self.ws_handler._get_shared_skill_manager()
|
||||
|
||||
# 如果启用了任务调度器,创建并启动
|
||||
if constants.SCHEDULER_ENABLED:
|
||||
self.scheduler = TaskScheduler(asyncio.get_running_loop(), self.logger)
|
||||
self.scheduler.set_websocket_handler(self.ws_handler)
|
||||
self.scheduler.set_ai_manager(self) # 设置AI管理器引用
|
||||
self.scheduler._get_shared_skill_manager()
|
||||
|
||||
async def handler(websocket):
|
||||
await self.handle_connection(websocket, None)
|
||||
@@ -2898,6 +3184,16 @@ class EnhancedAIManager:
|
||||
# Ping 帧,20s 内收不到 Pong 即强制断连。经 WAF→gateway 两层代理转发后 Pong 无法
|
||||
# 可靠回传,导致连接约 20s 被服务端关闭。前端已通过 25s 应用层心跳({type:'ping'})
|
||||
# 保活,故此处完全禁用服务端协议层 ping。
|
||||
# T1(审计 2026-08-21,产品决策:仅告警不加鉴权):绑定非回环地址时
|
||||
# WebSocket 端点(经工具链可达任意命令执行)即暴露给网络内任意
|
||||
# 客户端,启动时显著告警提醒确认网关/WAF 接入与网络边界
|
||||
_h = (host or "").strip().lower()
|
||||
if _h not in ("", "localhost", "127.0.0.1", "::1", "[::1]"):
|
||||
for _ in range(3):
|
||||
self.logger.warning(
|
||||
f"⚠️ WebSocket 绑定在 {host}:{port}(非回环地址):任何能建立 "
|
||||
f"TCP 连接的客户端都可使用本智能体(含命令执行能力),"
|
||||
f"请确认已置于可信网络或有网关/WAF 认证保护")
|
||||
server = await websockets.serve(handler, host, port, ping_interval=None)
|
||||
self.logger.info(i18n.get('websocket_started', host, port))
|
||||
|
||||
@@ -2914,6 +3210,8 @@ class EnhancedAIManager:
|
||||
finally:
|
||||
if self.scheduler:
|
||||
self.scheduler.stop()
|
||||
# I 批:释放命令专用线程池
|
||||
self.ws_handler.shutdown()
|
||||
|
||||
def run(self):
|
||||
"""运行WebSocket服务器"""
|
||||
@@ -2948,10 +3246,24 @@ class EnhancedAIManager:
|
||||
print(f"\n{i18n.get('error_prefix')}{e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
# W6(审计 2026-08-20):启动/运行失败此前打印堆栈后以退出码 0
|
||||
# 收场,supervisor/脚本误判成功
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
# ==================== 命令行参数处理 ====================
|
||||
|
||||
def _port_number(value):
|
||||
"""W6(审计 2026-08-20):--port argparse 类型校验,非法值启动前即报错"""
|
||||
try:
|
||||
iv = int(value)
|
||||
except ValueError:
|
||||
raise argparse.ArgumentTypeError(f"端口必须是整数: {value!r}")
|
||||
if not (1 <= iv <= 65535):
|
||||
raise argparse.ArgumentTypeError(f"端口必须在 1~65535 之间: {iv}")
|
||||
return iv
|
||||
|
||||
|
||||
def parse_args():
|
||||
"""解析命令行参数"""
|
||||
parser = argparse.ArgumentParser(description='单AI智能体系统 - WebSocket版本(所有任务都是对话任务)')
|
||||
@@ -2972,7 +3284,7 @@ def parse_args():
|
||||
|
||||
# WebSocket选项
|
||||
parser.add_argument('--host', type=str, help='WebSocket服务器主机')
|
||||
parser.add_argument('--port', type=int, help='WebSocket服务器端口')
|
||||
parser.add_argument('--port', type=_port_number, help='WebSocket服务器端口')
|
||||
|
||||
# 调度器选项
|
||||
parser.add_argument('--no-scheduler', action='store_true', help='禁用任务调度器')
|
||||
@@ -3056,9 +3368,11 @@ def handle_memo_commands(manager):
|
||||
if args.memo_add:
|
||||
try:
|
||||
parts = args.memo_add.split('|')
|
||||
title = parts[0]
|
||||
# W12(审计 2026-08-20):"|x" 空标题此前直接入库;"a|b|" 尾部空段
|
||||
# 以 "" 传入绕过 J1 时间校验——空串规整为 None,空标题由 db 层拦
|
||||
title = parts[0].strip()
|
||||
content = parts[1] if len(parts) > 1 else ""
|
||||
reminder = parts[2] if len(parts) > 2 else None
|
||||
reminder = (parts[2].strip() or None) if len(parts) > 2 else None
|
||||
|
||||
# 解析重复参数
|
||||
repeat_type = parts[3] if len(parts) > 3 else "none"
|
||||
@@ -3104,7 +3418,7 @@ def handle_memo_commands(manager):
|
||||
def handle_once_mode(manager, args) -> bool:
|
||||
"""一次性对话模式:命中 --once 则执行单次 AI 对话并以退出码结束进程。
|
||||
|
||||
设计要点(详见 docs/reports/aibox_cli_once_mode_dev_doc.md):
|
||||
设计要点:
|
||||
- stdout 仅输出 AI 回答正文(chunk 流式),line/reasoning 过程行走 stderr,管道安全
|
||||
- 退出码:0 成功 / 2 API 或链路失败 / 3 用户中断(Ctrl+C) / 4 全局超时
|
||||
- 复用三层中断机制(Event + HTTP 关闭 + 子进程 kill):SIGINT 与超时 Timer 均走 cancel()
|
||||
@@ -3186,6 +3500,19 @@ def handle_once_mode(manager, args) -> bool:
|
||||
def _on_timeout():
|
||||
state["timed_out"] = True
|
||||
manager.cancel()
|
||||
# I 批(F9,审计 2026-08-18):软超时兜底——cancel 后若 AI 卡死在
|
||||
# 不可中断路径(HTTP 请求空窗、忽略信号的子进程),think 可能永不
|
||||
# 返回、进程挂起(退出码 4 永远到不了)。宽限一个超时周期后强杀进程,
|
||||
# 保证 --once 有硬性结束语义
|
||||
def _hard_exit():
|
||||
print(f"\n⏰ 一次性对话强制超时({args.once_timeout * 2} 秒),已终止",
|
||||
file=sys.stderr)
|
||||
sys.stderr.flush()
|
||||
os._exit(4)
|
||||
|
||||
hard_timer = threading.Timer(max(1, args.once_timeout), _hard_exit)
|
||||
hard_timer.daemon = True
|
||||
hard_timer.start()
|
||||
timer = threading.Timer(max(1, args.once_timeout), _on_timeout)
|
||||
timer.daemon = True
|
||||
timer.start()
|
||||
@@ -3196,9 +3523,9 @@ def handle_once_mode(manager, args) -> bool:
|
||||
if state["exit_code"] == 0 and not state["interrupted"] and not state["timed_out"]:
|
||||
if response is None or (isinstance(response, str) and not response.strip()):
|
||||
state["exit_code"] = 2
|
||||
except KeyboardInterrupt:
|
||||
state["interrupted"] = True
|
||||
except Exception as e:
|
||||
# J15(审计 2026-08-19):except KeyboardInterrupt 为死代码——SIGINT
|
||||
# 已由自定义 handler 接管(_sigint_handler),KeyboardInterrupt 永不抛出
|
||||
print(f"{i18n.get('error_prefix')}{e}", file=sys.stderr)
|
||||
state["exit_code"] = 2
|
||||
finally:
|
||||
@@ -3238,6 +3565,17 @@ def main():
|
||||
|
||||
config = ConfigManager()
|
||||
|
||||
# W8(审计 2026-08-20):--temperature 入口钳制,防 99/-1 直达 API 400
|
||||
if args.temperature is not None:
|
||||
args.temperature = min(2.0, max(0.0, args.temperature))
|
||||
|
||||
# W11(审计 2026-08-20):config.json 的 system/scheduler 节被写成非 dict
|
||||
# 时,下方 CLI 覆盖赋值 TypeError 崩启动——回退默认节
|
||||
_section_defaults = config._get_default_config()
|
||||
for _section in ("system", "scheduler"):
|
||||
if not isinstance(config._config.get(_section), dict):
|
||||
config._config[_section] = _section_defaults.get(_section, {})
|
||||
|
||||
if args.debug is not None:
|
||||
config._config["system"]["debug_level"] = args.debug
|
||||
|
||||
@@ -3253,8 +3591,25 @@ def main():
|
||||
# 调度器设置
|
||||
if args.no_scheduler:
|
||||
config._config["scheduler"]["enabled"] = False
|
||||
if args.scheduler_interval:
|
||||
config._config["scheduler"]["check_interval"] = args.scheduler_interval
|
||||
# W7(审计 2026-08-20):负值直达 asyncio.sleep(负) 空转热循环,0 被 falsy
|
||||
# 静默忽略——统一钳制下限 0.1s
|
||||
config._config["scheduler"]["check_interval"] = max(0.1, args.scheduler_interval)
|
||||
|
||||
# W5(审计 2026-08-20):暂存 CLI 进程级覆盖,WS /reload 后重放(否则
|
||||
# config.reload() 重建 _config 会静默丢弃)
|
||||
_cli_overrides = {"system": {}, "scheduler": {}}
|
||||
if args.debug is not None:
|
||||
_cli_overrides["system"]["debug_level"] = args.debug
|
||||
if args.verbose:
|
||||
_cli_overrides["system"]["debug_level"] = 3
|
||||
if args.host:
|
||||
_cli_overrides["system"]["websocket_host"] = args.host
|
||||
if args.port:
|
||||
_cli_overrides["system"]["websocket_port"] = args.port
|
||||
if args.no_scheduler:
|
||||
_cli_overrides["scheduler"]["enabled"] = False
|
||||
_cli_overrides["scheduler"]["check_interval"] = config._config["scheduler"]["check_interval"]
|
||||
config._cli_overrides = {k: v for k, v in _cli_overrides.items() if v}
|
||||
|
||||
# 确保 ~/.aibox 目录存在
|
||||
os.makedirs(constants.SAVE_DIR, exist_ok=True)
|
||||
|
||||
@@ -815,6 +815,14 @@ async def main():
|
||||
await cli.run()
|
||||
except KeyboardInterrupt:
|
||||
print(f"\n{Colors.SYSTEM}程序已退出{Colors.RESET}")
|
||||
finally:
|
||||
# M8(审计 2026-08-20):正常退出(EOF 等)此前不关 websocket,
|
||||
# 服务端见异常断连——统一 finally 通知
|
||||
if cli.client:
|
||||
try:
|
||||
await cli.client.disconnect()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
common.atomic - 原子写文件工具(J5,审计 2026-08-19)
|
||||
|
||||
同目录 mkstemp 临时文件 + fsync + os.replace:进程中断/断电不留半截文件;
|
||||
mkstemp 随机文件名避免多进程并发写同路径互踩(conversation_history 旧实现
|
||||
的固定 .tmp 名在并发场景会互相覆盖)。仅依赖标准库,供 common/subsystems
|
||||
各写路径复用。
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
|
||||
|
||||
def atomic_write_text(path: str, content: str) -> None:
|
||||
"""原子写文本文件。失败时清理临时文件并抛出原异常。"""
|
||||
path = os.path.expanduser(path)
|
||||
directory = os.path.dirname(path) or "."
|
||||
os.makedirs(directory, exist_ok=True)
|
||||
fd, tmp_path = tempfile.mkstemp(prefix=".aibox_tmp_", dir=directory)
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
||||
f.write(content)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
# T27(审计 2026-08-21):Windows 下目标被其他进程打开(杀毒扫描/
|
||||
# 编辑器/并发读者)时 os.replace 抛 PermissionError。POSIX 无条件
|
||||
# 成功不受影响;Windows 做短退避重试,高频写路径(每条消息都原子
|
||||
# 写 current_conversation.json)一次冲突不再直接丢该条持久化。
|
||||
last_exc = None
|
||||
for attempt in range(3):
|
||||
try:
|
||||
os.replace(tmp_path, path)
|
||||
last_exc = None
|
||||
break
|
||||
except PermissionError as e:
|
||||
last_exc = e
|
||||
if attempt < 2:
|
||||
time.sleep(0.05 * (attempt + 1))
|
||||
if last_exc is not None:
|
||||
raise last_exc
|
||||
except BaseException:
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
|
||||
+50
-22
@@ -9,8 +9,12 @@ common.config - aibox 配置管理器
|
||||
import os
|
||||
import json
|
||||
import copy
|
||||
import shutil
|
||||
import threading
|
||||
from datetime import datetime
|
||||
from typing import Dict, Optional
|
||||
|
||||
from .atomic import atomic_write_text
|
||||
from .i18n import i18n
|
||||
class ConfigManager:
|
||||
"""配置管理器 - 单例模式,负责加载和保存配置"""
|
||||
@@ -18,15 +22,20 @@ class ConfigManager:
|
||||
_instance = None
|
||||
_config = None
|
||||
_config_path = None
|
||||
_instance_lock = threading.Lock()
|
||||
|
||||
def __new__(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
# J6(审计 2026-08-19):此前单例创建与初始化均无锁——多线程首例
|
||||
# 并发可双实例、双重加载配置
|
||||
with cls._instance_lock:
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
return cls._instance
|
||||
|
||||
def __init__(self):
|
||||
if self._config is None:
|
||||
self._load_config()
|
||||
with self._instance_lock:
|
||||
if self._config is None:
|
||||
self._load_config()
|
||||
|
||||
def _get_default_config_path(self) -> str:
|
||||
env_config = os.environ.get("AI_CONFIG_PATH")
|
||||
@@ -183,6 +192,10 @@ class ConfigManager:
|
||||
print(i18n.get('config_loaded', config_path))
|
||||
except Exception as e:
|
||||
print(f"⚠️ {i18n.get('error_prefix')}{e},使用默认配置")
|
||||
# J5(审计 2026-08-19):损坏文件先备份再继续——此前直接改用
|
||||
# 默认配置,任一后续保存会用"默认+本次改动"整体覆盖原文件,
|
||||
# proxy/api/context 等全部手工配置永久丢失
|
||||
self._backup_corrupted_config(config_path)
|
||||
self._config = self._get_default_config()
|
||||
else:
|
||||
print(i18n.get('config_not_found', config_path))
|
||||
@@ -190,6 +203,15 @@ class ConfigManager:
|
||||
self._save_default_config(config_path)
|
||||
self._migrate_legacy_proxy_password()
|
||||
|
||||
def _backup_corrupted_config(self, config_path: str):
|
||||
"""J5:损坏配置先复制备份(带时间戳),防止后续保存整体覆盖丢失"""
|
||||
try:
|
||||
backup = f"{config_path}.corrupt.{datetime.now().strftime('%Y%m%d_%H%M%S')}"
|
||||
shutil.copy2(config_path, backup)
|
||||
print(f"⚠️ 已备份损坏的配置文件到 {backup}")
|
||||
except Exception as e:
|
||||
print(f"⚠️ 备份损坏配置文件失败: {e}")
|
||||
|
||||
def _merge_configs(self, default: Dict, user: Dict) -> Dict:
|
||||
result = default.copy()
|
||||
for key, value in user.items():
|
||||
@@ -201,9 +223,9 @@ class ConfigManager:
|
||||
|
||||
def _save_default_config(self, config_path: str):
|
||||
try:
|
||||
os.makedirs(os.path.dirname(config_path), exist_ok=True)
|
||||
with open(config_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(self._get_default_config(), f, ensure_ascii=False, indent=2)
|
||||
atomic_write_text(
|
||||
config_path,
|
||||
json.dumps(self._get_default_config(), ensure_ascii=False, indent=2))
|
||||
print(i18n.get('config_created', config_path))
|
||||
except:
|
||||
pass
|
||||
@@ -278,27 +300,33 @@ class ConfigManager:
|
||||
"password_env": password_env
|
||||
}
|
||||
if save:
|
||||
self._save_config()
|
||||
return self._save_config()
|
||||
return True
|
||||
|
||||
def _save_config(self):
|
||||
"""将当前配置保存到配置文件(覆盖)。
|
||||
def _save_config(self) -> bool:
|
||||
"""将当前配置保存到配置文件(覆盖)。返回是否成功落盘。
|
||||
|
||||
落盘前对代理密码做纵深防御脱敏:无论调用路径如何,明文密码一律不写入文件。
|
||||
L2(审计 2026-08-18):此前异常仅 print,set_proxy 无条件报"已保存",
|
||||
磁盘只读等场景下用户以为已持久化,重启后配置丢失。
|
||||
"""
|
||||
config_path = self._config_path
|
||||
if config_path:
|
||||
try:
|
||||
safe_config = copy.deepcopy(self._config)
|
||||
socks = safe_config.get("proxy", {}).get("socks5")
|
||||
if socks and socks.get("password"):
|
||||
socks["password"] = None
|
||||
os.makedirs(os.path.dirname(config_path), exist_ok=True)
|
||||
with open(config_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(safe_config, f, ensure_ascii=False, indent=2)
|
||||
print(f"✅ 配置已保存到 {config_path}")
|
||||
except Exception as e:
|
||||
print(f"⚠️ 保存配置失败: {e}")
|
||||
if not config_path:
|
||||
return False
|
||||
try:
|
||||
safe_config = copy.deepcopy(self._config)
|
||||
socks = safe_config.get("proxy", {}).get("socks5")
|
||||
if socks and socks.get("password"):
|
||||
socks["password"] = None
|
||||
# J5:原子写(mkstemp + fsync + os.replace),崩溃不留半截 JSON
|
||||
atomic_write_text(
|
||||
config_path,
|
||||
json.dumps(safe_config, ensure_ascii=False, indent=2))
|
||||
print(f"✅ 配置已保存到 {config_path}")
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"⚠️ 保存配置失败: {e}")
|
||||
return False
|
||||
|
||||
def get_proxy_info(self) -> Dict:
|
||||
"""获取代理配置的详细信息(用于显示,不回传任何密码值)"""
|
||||
|
||||
@@ -236,6 +236,31 @@ class Constants:
|
||||
def COMPRESSOR_LLM_MAX_TOKENS(cls):
|
||||
return cls._get_config().get("compressor.llm_max_tokens", 1000)
|
||||
|
||||
# 后台进程管理(bg_run/bg_list/bg_output/bg_kill)
|
||||
@property
|
||||
def BG_PROCESSES_ENABLED(cls):
|
||||
return cls._get_config().get("background_processes.enabled", True)
|
||||
|
||||
@property
|
||||
def BG_MAX_PROCESSES(cls):
|
||||
return cls._get_config().get("background_processes.max_processes", 16)
|
||||
|
||||
@property
|
||||
def BG_MAX_ENTRIES(cls):
|
||||
return cls._get_config().get("background_processes.max_entries", 32)
|
||||
|
||||
@property
|
||||
def BG_BUFFER_PER_PROCESS(cls):
|
||||
return cls._get_config().get("background_processes.buffer_per_process", 262144)
|
||||
|
||||
@property
|
||||
def BG_MAX_LIFETIME(cls):
|
||||
return cls._get_config().get("background_processes.max_lifetime", 0)
|
||||
|
||||
@property
|
||||
def BG_MAX_WRITE_BYTES(cls):
|
||||
return cls._get_config().get("background_processes.max_write_bytes", 65536)
|
||||
|
||||
|
||||
# 创建常量实例
|
||||
constants = Constants()
|
||||
@@ -98,6 +98,7 @@ class I18n:
|
||||
'memo_overdue': "⚠️",
|
||||
'memo_pending': "⏰",
|
||||
'search_results': "🔍 搜索 '{}' 找到 {} 条记忆:",
|
||||
'search_usage': "用法: /search 关键词",
|
||||
'current_config': "⚙️ 当前配置:",
|
||||
'config_debug': " 调试级别: {}",
|
||||
'config_data_dir': " 数据目录: {}",
|
||||
@@ -110,6 +111,7 @@ class I18n:
|
||||
'enabled': "启用",
|
||||
'disabled': "禁用",
|
||||
'config_reloaded': "✅ 配置已重新加载",
|
||||
'config_reloaded_cli_kept': "(命令行进程级覆盖已保留)",
|
||||
'prompts_reloaded': "✅ 提示词已重新加载 ({} 个)",
|
||||
'checking_memos': "🔍 检查任务...",
|
||||
'overdue_found': "发现 {} 个到期任务:",
|
||||
@@ -157,6 +159,7 @@ class I18n:
|
||||
'memo_add_failed': "❌ 添加失败: {}",
|
||||
'memo_completed': "✅ 任务 #{} 已完成",
|
||||
'memo_complete_failed': "❌ 无法完成任务 #{}",
|
||||
'memo_complete_missing_id': "❌ 缺少 task_id,无法完成任务",
|
||||
|
||||
# 即时任务相关
|
||||
'immediate_memo_added': "⚡ 已添加为即时任务并立即检查",
|
||||
@@ -258,6 +261,7 @@ class I18n:
|
||||
'memo_overdue': "⚠️",
|
||||
'memo_pending': "⏰",
|
||||
'search_results': "🔍 Searching '{}' found {} memories:",
|
||||
'search_usage': "Usage: /search <keyword>",
|
||||
'current_config': "⚙️ Current configuration:",
|
||||
'config_debug': " Debug level: {}",
|
||||
'config_data_dir': " Data directory: {}",
|
||||
@@ -270,6 +274,7 @@ class I18n:
|
||||
'enabled': "enabled",
|
||||
'disabled': "disabled",
|
||||
'config_reloaded': "✅ Configuration reloaded",
|
||||
'config_reloaded_cli_kept': " (CLI process-level overrides preserved)",
|
||||
'prompts_reloaded': "✅ Prompts reloaded ({})",
|
||||
'checking_memos': "🔍 Checking tasks...",
|
||||
'overdue_found': "Found {} overdue tasks:",
|
||||
@@ -317,6 +322,7 @@ class I18n:
|
||||
'memo_add_failed': "❌ Failed to add: {}",
|
||||
'memo_completed': "✅ Task #{} completed",
|
||||
'memo_complete_failed': "❌ Cannot complete task #{}",
|
||||
'memo_complete_missing_id': "❌ Missing task_id, cannot complete task",
|
||||
|
||||
# Immediate task related
|
||||
'immediate_memo_added': "⚡ Added as immediate task and checked now",
|
||||
|
||||
+30
-12
@@ -4,12 +4,14 @@
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>AI智能助手 - Socket.IO客户端</title>
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css">
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/styles/github-dark.min.css">
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/marked/9.1.6/marked.min.js"></script>
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/highlight.min.js"></script>
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css" integrity="sha384-3B6NwesSXE7YJlcLI9RpRqGf2p/EgVH8BgoKTaUrmKNDkHPStTQ3EyoYjCGXaOTS" crossorigin="anonymous">
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/styles/github-dark.min.css" integrity="sha384-wH75j6z1lH97ZOpMOInqhgKzFkAInZPPSPlZpYKYTOqsaizPvhQZmAtLcPKXpLyH" crossorigin="anonymous">
|
||||
<!-- J13(审计 2026-08-19):CDN 脚本补 SRI integrity,防供应链劫持 -->
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/marked/9.1.6/marked.min.js" integrity="sha384-odPBjvtXVM/5hOYIr3A1dB+flh0c3wAT3bSesIOqEGmyUA4JoKf/YTWy0XKOYAY7" crossorigin="anonymous"></script>
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/highlight.min.js" integrity="sha384-g4mRvs7AO0/Ol5LxcGyz4Doe21pVhGNnC3EQw5shw+z+aXDN86HqUdwXWO+Gz2zI" crossorigin="anonymous"></script>
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/dompurify/3.2.4/purify.min.js" integrity="sha384-eEu5CTj3qGvu9PdJuS+YlkNi7d2XxQROAFYOr59zgObtlcux1ae1Il3u7jvdCSWu" crossorigin="anonymous"></script>
|
||||
<!-- Socket.IO 客户端库 -->
|
||||
<script src="https://cdn.socket.io/4.5.0/socket.io.min.js"></script>
|
||||
<script src="https://cdn.socket.io/4.5.0/socket.io.min.js" integrity="sha384-7EyYLQZgWBi67fBtVxw60/OWl1kjsfrPFcaU0pp0nAh+i8FD068QogUvg85Ewy1k" crossorigin="anonymous"></script>
|
||||
<style>
|
||||
/* 样式保持不变,完全使用原有的样式 */
|
||||
* {
|
||||
@@ -982,6 +984,15 @@
|
||||
</template>
|
||||
|
||||
<script>
|
||||
// J2(审计 2026-08-19):marked 输出经 DOMPurify 净化——AI 输出/命令
|
||||
// 执行结果是不可信输入,含原始 HTML 即可注入脚本;DOMPurify 缺失时
|
||||
// 退化为转义文本(只显示不执行,防 XSS 兜底)
|
||||
function renderMarkdown(content) {
|
||||
const html = marked.parse(content);
|
||||
if (window.DOMPurify) return DOMPurify.sanitize(html);
|
||||
return html.replace(/</g, '<');
|
||||
}
|
||||
|
||||
document.addEventListener('DOMContentLoaded', function() {
|
||||
// 配置marked
|
||||
marked.setOptions({
|
||||
@@ -1124,7 +1135,7 @@
|
||||
textDiv.setAttribute('data-raw-content', streamData.content);
|
||||
|
||||
if (markdownEnabled) {
|
||||
textDiv.innerHTML = marked.parse(streamData.content);
|
||||
textDiv.innerHTML = renderMarkdown(streamData.content);
|
||||
textDiv.querySelectorAll('pre code').forEach((block) => {
|
||||
hljs.highlightElement(block);
|
||||
});
|
||||
@@ -1419,7 +1430,7 @@
|
||||
const rawContent = msgDiv.getAttribute('data-raw-content');
|
||||
if (rawContent) {
|
||||
if (markdownEnabled) {
|
||||
msgDiv.innerHTML = marked.parse(rawContent);
|
||||
msgDiv.innerHTML = renderMarkdown(rawContent);
|
||||
msgDiv.classList.add('markdown-body');
|
||||
addCopyButtonsToCodeBlocks(msgDiv);
|
||||
} else {
|
||||
@@ -1438,6 +1449,7 @@
|
||||
let reconnectAttempts = 0;
|
||||
const MAX_RECONNECT_ATTEMPTS = 5;
|
||||
let targetWsUrl = null;
|
||||
let reconnectTimer = null; // F2(审计 2026-08-20):重连定时器句柄,手动断开时清除
|
||||
|
||||
function createSocketConnection() {
|
||||
targetWsUrl = document.getElementById('ws-url').value.trim();
|
||||
@@ -1503,7 +1515,8 @@
|
||||
if (reason !== 'io client disconnect' && reconnectAttempts < MAX_RECONNECT_ATTEMPTS) {
|
||||
reconnectAttempts++;
|
||||
showSystemMessage(`⚠️ 连接断开,${reconnectAttempts}/${MAX_RECONNECT_ATTEMPTS} 次重试...`);
|
||||
setTimeout(() => {
|
||||
reconnectTimer = setTimeout(() => {
|
||||
reconnectTimer = null;
|
||||
if (!socket || !socket.connected) {
|
||||
console.log('尝试重连...');
|
||||
socket = createSocketConnection();
|
||||
@@ -1676,7 +1689,8 @@
|
||||
|
||||
// 事件绑定
|
||||
document.getElementById('connect-btn').addEventListener('click', function() {
|
||||
if (socket && socket.connected) {
|
||||
// F3(审计 2026-08-20):CONNECTING 中的旧连接也要断开,否则重复点击产生孤儿连接/双后端会话
|
||||
if (socket) {
|
||||
socket.disconnect();
|
||||
}
|
||||
reconnectAttempts = 0;
|
||||
@@ -1684,6 +1698,8 @@
|
||||
});
|
||||
|
||||
document.getElementById('disconnect-btn').addEventListener('click', function() {
|
||||
// F2:清除已排程的重连定时器,防止倒计时到期把连接"复活"
|
||||
if (reconnectTimer) { clearTimeout(reconnectTimer); reconnectTimer = null; }
|
||||
if (socket) {
|
||||
socket.disconnect();
|
||||
socket = null;
|
||||
@@ -1717,10 +1733,12 @@
|
||||
|
||||
document.getElementById('clear-chat-btn').addEventListener('click', function() {
|
||||
const container = document.getElementById('messages-container');
|
||||
// F1(审计 2026-08-20):welcome-message 在 container 内,
|
||||
// 必须先 cloneNode 缓存再清空,否则 querySelector 返回 null 抛 TypeError
|
||||
const welcomeEl = document.querySelector('.welcome-message');
|
||||
const welcomeHtml = welcomeEl ? welcomeEl.cloneNode(true) : null;
|
||||
container.innerHTML = '';
|
||||
|
||||
const welcomeHtml = document.querySelector('.welcome-message').cloneNode(true);
|
||||
container.appendChild(welcomeHtml);
|
||||
if (welcomeHtml) container.appendChild(welcomeHtml);
|
||||
|
||||
clearAllStreamingMessages();
|
||||
});
|
||||
|
||||
+35
-11
@@ -5,13 +5,15 @@
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0, user-scalable=no">
|
||||
<title>小猫AI · 智能喵语伙伴</title>
|
||||
<!-- Font Awesome 6 (免费图标库) -->
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css">
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css" integrity="sha384-3B6NwesSXE7YJlcLI9RpRqGf2p/EgVH8BgoKTaUrmKNDkHPStTQ3EyoYjCGXaOTS" crossorigin="anonymous">
|
||||
<!-- highlight.js 代码高亮 + marked markdown解析 -->
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/styles/github-dark.min.css">
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/marked/9.1.6/marked.min.js"></script>
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/highlight.min.js"></script>
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/styles/github-dark.min.css" integrity="sha384-wH75j6z1lH97ZOpMOInqhgKzFkAInZPPSPlZpYKYTOqsaizPvhQZmAtLcPKXpLyH" crossorigin="anonymous">
|
||||
<!-- J13(审计 2026-08-19):CDN 脚本补 SRI integrity,防供应链劫持 -->
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/marked/9.1.6/marked.min.js" integrity="sha384-odPBjvtXVM/5hOYIr3A1dB+flh0c3wAT3bSesIOqEGmyUA4JoKf/YTWy0XKOYAY7" crossorigin="anonymous"></script>
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/highlight.min.js" integrity="sha384-g4mRvs7AO0/Ol5LxcGyz4Doe21pVhGNnC3EQw5shw+z+aXDN86HqUdwXWO+Gz2zI" crossorigin="anonymous"></script>
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/dompurify/3.2.4/purify.min.js" integrity="sha384-eEu5CTj3qGvu9PdJuS+YlkNi7d2XxQROAFYOr59zgObtlcux1ae1Il3u7jvdCSWu" crossorigin="anonymous"></script>
|
||||
<!-- Socket.IO 客户端库 -->
|
||||
<script src="https://cdn.socket.io/4.5.0/socket.io.min.js"></script>
|
||||
<script src="https://cdn.socket.io/4.5.0/socket.io.min.js" integrity="sha384-7EyYLQZgWBi67fBtVxw60/OWl1kjsfrPFcaU0pp0nAh+i8FD068QogUvg85Ewy1k" crossorigin="anonymous"></script>
|
||||
<style>
|
||||
/* ========== 融合小猫风格 + WebSocket增强样式 ========== */
|
||||
* {
|
||||
@@ -498,6 +500,15 @@
|
||||
|
||||
<script>
|
||||
(function(){
|
||||
// J2(审计 2026-08-19):marked 输出经 DOMPurify 净化——AI 输出/命令
|
||||
// 执行结果是不可信输入,含原始 HTML 即可注入脚本;DOMPurify 缺失时
|
||||
// 退化为转义文本(只显示不执行,防 XSS 兜底)
|
||||
function renderMarkdown(content) {
|
||||
const html = marked.parse(content);
|
||||
if (window.DOMPurify) return DOMPurify.sanitize(html);
|
||||
return html.replace(/</g, '<');
|
||||
}
|
||||
|
||||
// ---------- DOM 元素 ----------
|
||||
const bgFile = document.getElementById('bgFile');
|
||||
const roleImg = document.getElementById('roleImg');
|
||||
@@ -533,7 +544,12 @@
|
||||
const div = document.createElement('div');
|
||||
div.className = isError ? 'error-msg' : 'system-msg';
|
||||
div.style.cssText = 'background:rgba(0,0,0,0.6); border-radius:16px; padding:6px 12px; margin:6px 0; font-size:12px;';
|
||||
div.innerHTML = `<i class="fas ${isError ? 'fa-exclamation-triangle' : 'fa-info-circle'}"></i> ${text}`;
|
||||
// J2:text 为服务端内容(命令执行结果/AI 输出),禁止拼 innerHTML——
|
||||
// 原实现 `${text}` 模板拼接使任意 HTML 注入执行;图标为静态字面量
|
||||
const icon = document.createElement('i');
|
||||
icon.className = `fas ${isError ? 'fa-exclamation-triangle' : 'fa-info-circle'}`;
|
||||
div.appendChild(icon);
|
||||
div.appendChild(document.createTextNode(' ' + text));
|
||||
chatBox.appendChild(div);
|
||||
chatBox.scrollTop = chatBox.scrollHeight;
|
||||
}
|
||||
@@ -545,7 +561,7 @@
|
||||
if(textDiv) {
|
||||
textDiv.setAttribute('data-raw', stream.content);
|
||||
if(markdownEnabled) {
|
||||
textDiv.innerHTML = marked.parse(stream.content);
|
||||
textDiv.innerHTML = renderMarkdown(stream.content);
|
||||
textDiv.querySelectorAll('pre code').forEach(block => hljs.highlightElement(block));
|
||||
} else textDiv.innerText = stream.content;
|
||||
}
|
||||
@@ -572,7 +588,12 @@
|
||||
if(!streamingTasks[eid]) {
|
||||
const sep = document.createElement('div');
|
||||
sep.className = 'task-separator';
|
||||
sep.innerHTML = `<i class="fas fa-bolt"></i> 定时任务 #${taskMeta.taskId} (${taskMeta.taskTitle || '执行'}) 第${taskMeta.triggerCount||1}次执行`;
|
||||
// J2:taskMeta.taskTitle 为服务端可控内容,禁止拼 innerHTML
|
||||
const bolt = document.createElement('i');
|
||||
bolt.className = 'fas fa-bolt';
|
||||
sep.appendChild(bolt);
|
||||
sep.appendChild(document.createTextNode(
|
||||
` 定时任务 #${taskMeta.taskId} (${taskMeta.taskTitle || '执行'}) 第${taskMeta.triggerCount||1}次执行`));
|
||||
chatBox.appendChild(sep);
|
||||
const taskDiv = document.createElement('div');
|
||||
taskDiv.className = 'task-msg streaming';
|
||||
@@ -590,7 +611,7 @@
|
||||
if(st && st.element) {
|
||||
const textDiv = st.element.querySelector('.message-text');
|
||||
if(textDiv) {
|
||||
if(markdownEnabled) textDiv.innerHTML = marked.parse(st.content);
|
||||
if(markdownEnabled) textDiv.innerHTML = renderMarkdown(st.content);
|
||||
else textDiv.innerText = st.content;
|
||||
textDiv.querySelectorAll('pre code').forEach(block => hljs.highlightElement(block));
|
||||
}
|
||||
@@ -613,7 +634,7 @@
|
||||
if(task.timer) clearTimeout(task.timer);
|
||||
if(task.element) {
|
||||
const txtDiv = task.element.querySelector('.message-text');
|
||||
if(txtDiv && markdownEnabled) txtDiv.innerHTML = marked.parse(task.content);
|
||||
if(txtDiv && markdownEnabled) txtDiv.innerHTML = renderMarkdown(task.content);
|
||||
task.element.classList.remove('streaming');
|
||||
}
|
||||
globalMsgCount++;
|
||||
@@ -652,7 +673,8 @@
|
||||
}
|
||||
|
||||
function connectWebSocket() {
|
||||
if(socket && socket.connected) {
|
||||
// F3(审计 2026-08-20):CONNECTING 中的旧连接也要断开,否则重复点击产生孤儿连接/双后端会话
|
||||
if(socket) {
|
||||
socket.disconnect();
|
||||
}
|
||||
|
||||
@@ -731,6 +753,8 @@
|
||||
|
||||
function disconnectWebSocket() {
|
||||
manualDisconnect = true;
|
||||
// F2(审计 2026-08-20):清除已排程的重连定时器,防止倒计时到期把连接"复活"
|
||||
if(reconnectTimer) { clearTimeout(reconnectTimer); reconnectTimer = null; }
|
||||
if(socket) {
|
||||
socket.disconnect();
|
||||
socket = null;
|
||||
|
||||
+30
-12
@@ -4,10 +4,12 @@
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>AI智能助手 - WebSocket直连</title>
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css">
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/styles/github-dark.min.css">
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/marked/9.1.6/marked.min.js"></script>
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/highlight.min.js"></script>
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css" integrity="sha384-3B6NwesSXE7YJlcLI9RpRqGf2p/EgVH8BgoKTaUrmKNDkHPStTQ3EyoYjCGXaOTS" crossorigin="anonymous">
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/styles/github-dark.min.css" integrity="sha384-wH75j6z1lH97ZOpMOInqhgKzFkAInZPPSPlZpYKYTOqsaizPvhQZmAtLcPKXpLyH" crossorigin="anonymous">
|
||||
<!-- J13(审计 2026-08-19):CDN 脚本补 SRI integrity,防供应链劫持 -->
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/marked/9.1.6/marked.min.js" integrity="sha384-odPBjvtXVM/5hOYIr3A1dB+flh0c3wAT3bSesIOqEGmyUA4JoKf/YTWy0XKOYAY7" crossorigin="anonymous"></script>
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/highlight.min.js" integrity="sha384-g4mRvs7AO0/Ol5LxcGyz4Doe21pVhGNnC3EQw5shw+z+aXDN86HqUdwXWO+Gz2zI" crossorigin="anonymous"></script>
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/dompurify/3.2.4/purify.min.js" integrity="sha384-eEu5CTj3qGvu9PdJuS+YlkNi7d2XxQROAFYOr59zgObtlcux1ae1Il3u7jvdCSWu" crossorigin="anonymous"></script>
|
||||
<style>
|
||||
/* 样式保持不变,完全使用原有的样式 */
|
||||
* {
|
||||
@@ -951,6 +953,15 @@
|
||||
</template>
|
||||
|
||||
<script>
|
||||
// J2(审计 2026-08-19):marked 输出经 DOMPurify 净化——AI 输出/命令
|
||||
// 执行结果是不可信输入,含原始 HTML 即可注入脚本;DOMPurify 缺失时
|
||||
// 退化为转义文本(只显示不执行,防 XSS 兜底)
|
||||
function renderMarkdown(content) {
|
||||
const html = marked.parse(content);
|
||||
if (window.DOMPurify) return DOMPurify.sanitize(html);
|
||||
return html.replace(/</g, '<');
|
||||
}
|
||||
|
||||
document.addEventListener('DOMContentLoaded', function() {
|
||||
// 配置marked
|
||||
marked.setOptions({
|
||||
@@ -1048,7 +1059,7 @@
|
||||
textDiv.setAttribute('data-raw-content', streamData.content);
|
||||
|
||||
if (markdownEnabled) {
|
||||
textDiv.innerHTML = marked.parse(streamData.content);
|
||||
textDiv.innerHTML = renderMarkdown(streamData.content);
|
||||
textDiv.querySelectorAll('pre code').forEach((block) => {
|
||||
hljs.highlightElement(block);
|
||||
});
|
||||
@@ -1336,7 +1347,7 @@
|
||||
const rawContent = msgDiv.getAttribute('data-raw-content');
|
||||
if (rawContent) {
|
||||
if (markdownEnabled) {
|
||||
msgDiv.innerHTML = marked.parse(rawContent);
|
||||
msgDiv.innerHTML = renderMarkdown(rawContent);
|
||||
msgDiv.classList.add('markdown-body');
|
||||
addCopyButtonsToCodeBlocks(msgDiv);
|
||||
} else {
|
||||
@@ -1413,9 +1424,12 @@
|
||||
|
||||
if (!manualDisconnect && reconnectAttempts < MAX_RECONNECT_ATTEMPTS) {
|
||||
reconnectAttempts++;
|
||||
showSystemMessage(`⚠️ 连接断开,${reconnectAttempts}/${MAX_RECONNECT_ATTEMPTS} 秒后重试...`);
|
||||
// F6(审计 2026-08-20):实际延迟为 3000*N ms,文案与此对齐
|
||||
showSystemMessage(`⚠️ 连接断开,${3 * reconnectAttempts} 秒后重试(${reconnectAttempts}/${MAX_RECONNECT_ATTEMPTS})...`);
|
||||
setTimeout(() => {
|
||||
if (!ws || ws.readyState !== WebSocket.OPEN) {
|
||||
// F2(审计 2026-08-20):闭包到期时再判 manualDisconnect,
|
||||
// 防止手动断开后已排程定时器把连接"复活"
|
||||
if (!manualDisconnect && (!ws || ws.readyState !== WebSocket.OPEN)) {
|
||||
console.log('尝试重连...');
|
||||
ws = connectWebSocket();
|
||||
}
|
||||
@@ -1585,7 +1599,8 @@
|
||||
|
||||
// 事件绑定
|
||||
document.getElementById('connect-btn').addEventListener('click', function() {
|
||||
if (ws && ws.readyState === WebSocket.OPEN) {
|
||||
// F3(审计 2026-08-20):CONNECTING 中的旧连接也要关闭,否则重复点击产生孤儿连接/双后端会话
|
||||
if (ws && ws.readyState !== WebSocket.CLOSED) {
|
||||
ws.close();
|
||||
}
|
||||
reconnectAttempts = 0;
|
||||
@@ -1628,10 +1643,13 @@
|
||||
|
||||
document.getElementById('clear-chat-btn').addEventListener('click', function() {
|
||||
const container = document.getElementById('messages-container');
|
||||
// F1(审计 2026-08-20):welcome-message 在 container 内,
|
||||
// 必须先 cloneNode 缓存再清空,否则 querySelector 返回 null 抛 TypeError,
|
||||
// 且后续 clearAllStreamingMessages/计数清零全部中断不执行
|
||||
const welcomeEl = document.querySelector('.welcome-message');
|
||||
const welcomeHtml = welcomeEl ? welcomeEl.cloneNode(true) : null;
|
||||
container.innerHTML = '';
|
||||
|
||||
const welcomeHtml = document.querySelector('.welcome-message').cloneNode(true);
|
||||
container.appendChild(welcomeHtml);
|
||||
if (welcomeHtml) container.appendChild(welcomeHtml);
|
||||
|
||||
clearAllStreamingMessages();
|
||||
messageCount = 0;
|
||||
|
||||
+102
-23
@@ -22,9 +22,36 @@ from typing import Dict, Iterator, List, Optional
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
import requests
|
||||
|
||||
# L12(审计 2026-08-18):瞬时错误体可能回显网关/WAF 错误页中的
|
||||
# Authorization 头,抹掉 Bearer 令牌后再进日志,防密钥落盘
|
||||
_AUTH_HEADER_RE = re.compile(r"Bearer\s+[A-Za-z0-9._~+/=-]+")
|
||||
# T13:错误响应体可能回显裸 API Key(如 sk-xxx),一并脱敏
|
||||
_API_KEY_RE = re.compile(r"sk-[A-Za-z0-9_-]{16,}")
|
||||
|
||||
# T19(审计 2026-08-21):瞬时网络异常白名单。此前 FailoverClient 捕获
|
||||
# requests.exceptions.RequestException 基类——档案 api_url 手写错(漏
|
||||
# scheme 抛 MissingSchema)等永久性配置错误也被当瞬时:标记冷却 60s、
|
||||
# 触发转移、计入 failovers,每次冷却到期再撞一次坏端点,统计与冷却表
|
||||
# 被永久错误污染。显式枚举瞬时异常,其余 requests 异常直接上抛响亮暴露
|
||||
_TRANSIENT_NET_ERRORS = (
|
||||
requests.exceptions.ConnectionError,
|
||||
requests.exceptions.Timeout,
|
||||
requests.exceptions.ChunkedEncodingError,
|
||||
)
|
||||
|
||||
|
||||
def _sanitize_log_msg(msg: str) -> str:
|
||||
"""日志脱敏:替换错误消息中的 Bearer 令牌与裸 API Key。"""
|
||||
if not msg:
|
||||
return msg
|
||||
msg = _AUTH_HEADER_RE.sub("Bearer ***", msg)
|
||||
msg = _API_KEY_RE.sub("sk-***", msg)
|
||||
return msg
|
||||
|
||||
|
||||
class LLMClient(ABC):
|
||||
"""大模型客户端抽象基类。子类需实现流式/非流式两种调用,并归一化为统一事件格式。"""
|
||||
@@ -64,7 +91,7 @@ class LLMClient(ABC):
|
||||
def _log(self, level: str, msg: str):
|
||||
if self.logger is not None and hasattr(self.logger, level):
|
||||
try:
|
||||
getattr(self.logger, level)(msg)
|
||||
getattr(self.logger, level)(_sanitize_log_msg(msg))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@@ -154,20 +181,29 @@ class DeepSeekClient(LLMClient):
|
||||
proxies=self._resolve_proxies(),
|
||||
)
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
return result["choices"][0]["message"]["content"].strip()
|
||||
if response.status_code == 429 or response.status_code >= 500:
|
||||
try:
|
||||
result = response.json()
|
||||
return result["choices"][0]["message"]["content"].strip()
|
||||
finally:
|
||||
response.close()
|
||||
# T9+T13:非 200 立即关连接(stream=True 才需要,但统一防御),
|
||||
# 错误体脱敏后再入异常,防 API Key 经错误页回显进历史/日志
|
||||
status = response.status_code
|
||||
body = _sanitize_log_msg(response.text[:200])
|
||||
response.close()
|
||||
if status == 429 or status >= 500:
|
||||
# L4a:瞬时故障显式抛出——FailoverClient 据此转移下一档案;
|
||||
# 无包装器时由调用方 catch-all 接住,下游仍走失败分支(语义等价)
|
||||
raise TransientLLMError(
|
||||
f"LLM 请求瞬时失败: HTTP {response.status_code}: {response.text[:200]}",
|
||||
f"LLM 请求瞬时失败: HTTP {status}: {body}",
|
||||
retry_after=_retry_after_of(response))
|
||||
self._log("error", f"LLM 压缩请求失败: HTTP {response.status_code}")
|
||||
self._log("error", f"LLM 压缩请求失败: HTTP {status}")
|
||||
return None
|
||||
except TransientLLMError:
|
||||
raise # L4a:瞬时故障上抛供 FailoverClient 转移,不吞
|
||||
except (requests.exceptions.ConnectionError, requests.exceptions.Timeout):
|
||||
raise # L4a:连接类瞬时故障同样上抛(原被吞成 None,无法转移)
|
||||
except _TRANSIENT_NET_ERRORS:
|
||||
raise # T19:连接类瞬时故障同样上抛(原被吞成 None,无法转移;
|
||||
# ChunkedEncodingError 一并对齐——响应体中途损坏属瞬时网络故障)
|
||||
except Exception as e:
|
||||
self._log("error", f"LLM 压缩请求异常: {e}")
|
||||
return None
|
||||
@@ -207,23 +243,32 @@ class DeepSeekClient(LLMClient):
|
||||
self._current_response = response
|
||||
|
||||
if response.status_code != 200:
|
||||
if response.status_code == 429 or response.status_code >= 500:
|
||||
# T9:非 200 立即释放连接(stream=True 下不关闭会阻塞同域名
|
||||
# in-flight 请求直到 GC);T13:错误体脱敏后再入异常
|
||||
status = response.status_code
|
||||
retry_after = _retry_after_of(response)
|
||||
body = _sanitize_log_msg(response.text[:200])
|
||||
response.close()
|
||||
self._current_response = None
|
||||
if status == 429 or status >= 500:
|
||||
# L4b:瞬时故障显式分类(此时尚未产出任何事件,可安全转移重放)
|
||||
raise TransientLLMError(
|
||||
f"LLM 流式请求瞬时失败: HTTP {response.status_code}: {response.text[:200]}",
|
||||
retry_after=_retry_after_of(response))
|
||||
f"LLM 流式请求瞬时失败: HTTP {status}: {body}",
|
||||
retry_after=retry_after)
|
||||
raise RuntimeError(
|
||||
f"HTTP {response.status_code}: {response.text[:200]}")
|
||||
f"HTTP {status}: {body}")
|
||||
|
||||
saw_done = False
|
||||
try:
|
||||
for line in response.iter_lines():
|
||||
if not line:
|
||||
continue
|
||||
line_text = line.decode("utf-8")
|
||||
line_text = line.decode("utf-8", errors="replace") # L3:非 UTF-8 行不中断流
|
||||
if not line_text.startswith("data: "):
|
||||
continue
|
||||
data = line_text[6:]
|
||||
if data == "[DONE]":
|
||||
saw_done = True
|
||||
yield {"type": "done"}
|
||||
break
|
||||
try:
|
||||
@@ -239,6 +284,14 @@ class DeepSeekClient(LLMClient):
|
||||
yield event
|
||||
finally:
|
||||
self._current_response = None
|
||||
# T20(审计 2026-08-21):iter_lines 耗尽但从未收到 [DONE](代理缓冲
|
||||
# 后提前 FIN / 服务端中途断开,末行半截 JSON 解析失败被 continue)
|
||||
# 此前被静默视为成功——回答无声截断无任何信号。抛瞬时异常:首事件
|
||||
# 前 FailoverClient 可转移重放;已产出事件则上抛(既有语义:已输出
|
||||
# 内容绝不重放),由上层错误处理回滚/报错。用户中断(cancel 关连接
|
||||
# 抛异常)不会走到这里。
|
||||
if not saw_done:
|
||||
raise TransientLLMError("LLM 流式响应不完整:未收到 [DONE] 结束标记(连接中途断开)")
|
||||
|
||||
# L5:思考内容的常见厂商字段别名(首个命中的非空字段生效)
|
||||
REASONING_FIELDS = ("reasoning_content", "reasoning", "thinking", "thinking_content")
|
||||
@@ -352,12 +405,21 @@ class FailoverClient:
|
||||
def _log(self, level: str, msg: str):
|
||||
if self.logger is not None and hasattr(self.logger, level):
|
||||
try:
|
||||
getattr(self.logger, level)(msg)
|
||||
getattr(self.logger, level)(_sanitize_log_msg(msg))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _cool_key(client) -> int:
|
||||
"""T6:冷却键按客户端实例(id)——client.name 取 display_name 优先且
|
||||
无唯一性校验,两个 profile_name 不同但 display_name 相同的档案
|
||||
(如主备都叫 "Kimi")冷却状态互相污染:A 失败冷却 → B 被误判
|
||||
跳过(该转移时不转移)。P6 去重保证每档案恰好一个实例,实例被
|
||||
_clients 持有生命周期内 id 稳定,与去重键语义一致。"""
|
||||
return id(client)
|
||||
|
||||
def _cooling(self, client) -> bool:
|
||||
entry = self._fail_at.get(client.name)
|
||||
entry = self._fail_at.get(self._cool_key(client))
|
||||
if entry is None:
|
||||
return False
|
||||
t, cd = entry
|
||||
@@ -365,9 +427,9 @@ class FailoverClient:
|
||||
|
||||
def _mark_fail(self, client, cooldown: Optional[float] = None):
|
||||
# L4c:429 Retry-After(经 TransientLLMError.retry_after)覆盖默认冷却时长
|
||||
self._fail_at[client.name] = (time.monotonic(),
|
||||
cooldown if cooldown and cooldown > 0
|
||||
else self.cooldown_seconds)
|
||||
self._fail_at[self._cool_key(client)] = (time.monotonic(),
|
||||
cooldown if cooldown and cooldown > 0
|
||||
else self.cooldown_seconds)
|
||||
|
||||
def _budget_exhausted(self) -> bool:
|
||||
if self.timeout_budget is None or self._call_started is None:
|
||||
@@ -382,7 +444,7 @@ class FailoverClient:
|
||||
self._stats["failovers"] += 1
|
||||
try:
|
||||
if self.on_failover is not None:
|
||||
self.on_failover(from_name, to_name, reason)
|
||||
self.on_failover(from_name, to_name, _sanitize_log_msg(reason))
|
||||
except Exception as e:
|
||||
self._log("warning", f"on_failover 回调异常(忽略): {e}")
|
||||
|
||||
@@ -422,7 +484,7 @@ class FailoverClient:
|
||||
return result
|
||||
# None:认证/参数类 4xx 或响应异常——不可转移,直接透传
|
||||
return None
|
||||
except (TransientLLMError, requests.exceptions.RequestException) as e:
|
||||
except (TransientLLMError,) + _TRANSIENT_NET_ERRORS as e:
|
||||
if self._cancelled:
|
||||
return None # 用户中断引发的异常:不转移
|
||||
last_exc = e
|
||||
@@ -439,10 +501,25 @@ class FailoverClient:
|
||||
max_tokens: int = 8192, tools: Optional[List[Dict]] = None,
|
||||
tool_choice: Optional[str] = None,
|
||||
reasoning_enabled: bool = False, **extra) -> Iterator[Dict]:
|
||||
"""L4b:流式转移——仅首事件前故障可安全重放(对用户不可见)。"""
|
||||
"""L4b:流式转移——仅首事件前故障可安全重放(对用户不可见)。
|
||||
|
||||
T7:本方法是普通函数(非生成器)——生成器函数体延迟到首次 next()
|
||||
才执行,若 cancel() 在创建生成器与首迭代之间被其他线程调用,
|
||||
延迟执行的 `_cancelled = False` 会把用户的取消清掉(请求照常发出)。
|
||||
标志重置与计数在返回生成器前同步完成,取消契约(类文档)恢复成立。
|
||||
"""
|
||||
self._cancelled = False # 新一轮请求,清除上一轮中断标志
|
||||
self._call_started = time.monotonic()
|
||||
self._stats["requests"] += 1
|
||||
return self._stream_chat_impl(messages, temperature=temperature,
|
||||
max_tokens=max_tokens, tools=tools,
|
||||
tool_choice=tool_choice,
|
||||
reasoning_enabled=reasoning_enabled, **extra)
|
||||
|
||||
def _stream_chat_impl(self, messages: List[Dict], temperature: float = None,
|
||||
max_tokens: int = 8192, tools: Optional[List[Dict]] = None,
|
||||
tool_choice: Optional[str] = None,
|
||||
reasoning_enabled: bool = False, **extra) -> Iterator[Dict]:
|
||||
last_exc = None
|
||||
for idx, client in enumerate(self._clients):
|
||||
if self._cooling(client):
|
||||
@@ -462,7 +539,7 @@ class FailoverClient:
|
||||
produced_any = True
|
||||
yield event
|
||||
return # 自然结束(含 done 事件)视为成功
|
||||
except (TransientLLMError, requests.exceptions.RequestException) as e:
|
||||
except (TransientLLMError,) + _TRANSIENT_NET_ERRORS as e:
|
||||
if self._cancelled:
|
||||
raise # 用户中断(cancel 强关连接):绝不转移
|
||||
if produced_any:
|
||||
@@ -515,7 +592,9 @@ def resolve_api_key(api_key: Optional[str], config=None,
|
||||
f"模型档案端点缺少 API Key:请设置环境变量 {env_name},"
|
||||
f"或将该档案的 api_key_env 配置为空串以启用本地免密端点")
|
||||
try:
|
||||
home = os.environ.get("HOME", os.environ["USERPROFILE"])
|
||||
# dict.get 的 default 参数会无条件先求值,直接写 os.environ["USERPROFILE"]
|
||||
# 在无该变量的 Linux 上必抛 KeyError,导致 key 文件回退路径永远不可达
|
||||
home = os.environ.get("HOME") or os.environ.get("USERPROFILE")
|
||||
key_path = os.path.join(home, ".cert", "deepseek.key")
|
||||
with open(key_path, "r") as f:
|
||||
return f.read().strip()
|
||||
|
||||
@@ -0,0 +1,492 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
大模型接入点(模型档案)交互式管理工具
|
||||
|
||||
不用手写 config.json:菜单式列出 / 添加 / 删除 / 连通性验证
|
||||
`models` 节的模型接入点(OpenAI Chat Completions 兼容端点:
|
||||
DeepSeek / Kimi / Qwen / GLM / 本地 Ollama 等)。
|
||||
|
||||
用法:
|
||||
python model_config.py # 交互式菜单
|
||||
python model_config.py --list # 非交互:列出当前接入点
|
||||
|
||||
安全约定(与 aibox 审计一致):API Key 一律不写入配置文件——
|
||||
只保存环境变量名(api_key_env),或留空表示免密本地端点。
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
# 允许从仓库根目录或 aibox/ 目录启动
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
from common import ConfigManager # noqa: E402
|
||||
from llm_client import _sanitize_log_msg # noqa: E402
|
||||
|
||||
PROFILE_NAME_RE = re.compile(r"^[A-Za-z0-9_-]+$")
|
||||
|
||||
# 常见厂商端点提示(首行是默认建议值)
|
||||
PROVIDER_HINTS = {
|
||||
"DeepSeek": "https://api.deepseek.com/v1/chat/completions",
|
||||
"Kimi (Moonshot)": "https://api.moonshot.cn/v1/chat/completions",
|
||||
"Qwen (阿里云百炼)": "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions",
|
||||
"GLM (智谱)": "https://open.bigmodel.cn/api/paas/v4/chat/completions",
|
||||
"Ollama 本地": "http://localhost:11434/v1/chat/completions",
|
||||
}
|
||||
|
||||
|
||||
# ==================== 核心逻辑(可单测,不依赖 input) ====================
|
||||
|
||||
def get_models(config: ConfigManager) -> Dict:
|
||||
"""取 models 节(不存在则初始化为空 dict)。"""
|
||||
models = config.get("models") or {}
|
||||
if not isinstance(models, dict):
|
||||
models = {}
|
||||
return models
|
||||
|
||||
|
||||
def _api_section(config: ConfigManager) -> Dict:
|
||||
"""M3(审计 2026-08-20):api 节被改坏为非 dict 时重置为空节——
|
||||
否则 setdefault 后索引赋值抛 TypeError/AttributeError"""
|
||||
api = config._config.get("api")
|
||||
if not isinstance(api, dict):
|
||||
api = config._config["api"] = {}
|
||||
return api
|
||||
|
||||
|
||||
def list_profiles(config: ConfigManager) -> List[Dict]:
|
||||
"""列出所有档案的归一化展示数据。"""
|
||||
models = get_models(config)
|
||||
default = config.get("api.default_model") or ""
|
||||
profiles = []
|
||||
for name, profile in models.items():
|
||||
if not isinstance(profile, dict):
|
||||
continue
|
||||
# M4(审计 2026-08-20):fallback 被改坏为字符串时原样下传,
|
||||
# 展示层 join 成单字符误导——统一规整为列表
|
||||
fallback = profile.get("fallback")
|
||||
profiles.append({
|
||||
"name": name,
|
||||
"is_default": name == default,
|
||||
"display_name": profile.get("display_name") or name,
|
||||
"api_url": profile.get("api_url", ""),
|
||||
"model": profile.get("model", ""),
|
||||
"api_key_env": profile.get("api_key_env", ""),
|
||||
"reasoning_enabled": bool(profile.get("reasoning_enabled", False)),
|
||||
"temperature": profile.get("temperature"),
|
||||
"timeout": profile.get("timeout"),
|
||||
"fallback": fallback if isinstance(fallback, list) else [],
|
||||
})
|
||||
return profiles
|
||||
|
||||
|
||||
def validate_new_profile(config: ConfigManager, name: str, api_url: str,
|
||||
model: str) -> Optional[str]:
|
||||
"""校验新档案字段,返回错误信息(合法返回 None)。"""
|
||||
if not name:
|
||||
return "档案名不能为空"
|
||||
if not PROFILE_NAME_RE.match(name):
|
||||
return f"档案名只能包含字母/数字/下划线/短横线(当前: {name})"
|
||||
models = get_models(config)
|
||||
if name in models:
|
||||
return f"档案名 '{name}' 已存在"
|
||||
if not api_url or not str(api_url).strip().lower().startswith(("http://", "https://")):
|
||||
return "API 地址必须是 http:// 或 https:// 开头的完整 URL"
|
||||
if not model or not str(model).strip():
|
||||
return "模型名不能为空(如 deepseek-chat / gpt-4o / qwen-plus)"
|
||||
return None
|
||||
|
||||
|
||||
def add_profile(config: ConfigManager, name: str, api_url: str, model: str,
|
||||
api_key_env: str = "", display_name: str = None,
|
||||
temperature: float = None, timeout: int = None,
|
||||
reasoning_enabled: bool = False,
|
||||
fallback: Optional[List[str]] = None,
|
||||
set_default: bool = False) -> List[str]:
|
||||
"""添加档案。返回提示信息列表;字段非法时返回带 ❌ 的错误信息。"""
|
||||
err = validate_new_profile(config, name, api_url, model)
|
||||
if err:
|
||||
return [f"❌ {err}"]
|
||||
|
||||
# M5(审计 2026-08-20):api_key_env 需为合法环境变量名;展示/地址字段
|
||||
# 拒控制字符并限长——防终端注入回显(\x1b] 等序列写入配置后被 print)
|
||||
api_key_env = (api_key_env or "").strip()
|
||||
if api_key_env and not re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", api_key_env):
|
||||
return [f"❌ 环境变量名不合法: {api_key_env!r}(须为字母/数字/下划线且不以数字开头)"]
|
||||
for label, val, limit in (("显示名", display_name, 64),
|
||||
("模型名", model, 128),
|
||||
("API 地址", api_url, 256)):
|
||||
if val and (len(str(val)) > limit or not all(c.isprintable() for c in str(val))):
|
||||
return [f"❌ {label}含控制字符或超过 {limit} 字符"]
|
||||
|
||||
models = get_models(config)
|
||||
profile = {
|
||||
"api_url": str(api_url).strip(),
|
||||
"model": str(model).strip(),
|
||||
}
|
||||
if api_key_env:
|
||||
profile["api_key_env"] = api_key_env.strip()
|
||||
else:
|
||||
# 空串 = 免密本地端点(Ollama/vLLM),显式声明
|
||||
profile["api_key_env"] = ""
|
||||
if display_name:
|
||||
profile["display_name"] = display_name.strip()
|
||||
if temperature is not None:
|
||||
try:
|
||||
profile["temperature"] = max(0.0, min(2.0, float(temperature)))
|
||||
except (TypeError, ValueError):
|
||||
return ["❌ 温度必须是 0.0~2.0 的数字"]
|
||||
if timeout is not None:
|
||||
try:
|
||||
profile["timeout"] = max(5, int(timeout))
|
||||
except (TypeError, ValueError):
|
||||
return ["❌ 超时必须是整数秒"]
|
||||
if reasoning_enabled:
|
||||
profile["reasoning_enabled"] = True
|
||||
if fallback:
|
||||
# 只保留已存在的档案,避免悬空引用(L4a 启动期硬性报错)
|
||||
known = set(models.keys())
|
||||
invalid = [f for f in fallback if f not in known or f == name]
|
||||
if invalid:
|
||||
return [f"❌ 失败转移包含不存在的档案或自身: {invalid},请重新选择"]
|
||||
profile["fallback"] = list(fallback)
|
||||
|
||||
models[name] = profile
|
||||
config._config["models"] = models
|
||||
if set_default:
|
||||
_api_section(config)["default_model"] = name
|
||||
|
||||
saved = config._save_config()
|
||||
msgs = [f"✅ 接入点 '{name}' 已添加({profile['api_url']},模型 {profile['model']})"]
|
||||
if not saved:
|
||||
msgs.append("⚠️ 保存失败(仅本次运行生效,重启后丢失)——请检查配置文件目录权限")
|
||||
if set_default:
|
||||
msgs.append(f"⭐ 已设为默认接入点(--model 不指定时使用)")
|
||||
if not api_key_env:
|
||||
msgs.append("ℹ️ 免密端点:请求不携带 API Key(适用于 Ollama/vLLM 等本地服务)")
|
||||
else:
|
||||
msgs.append(f"🔑 Key 从环境变量 {api_key_env} 读取(未设置时该档案不可用)")
|
||||
return msgs
|
||||
|
||||
|
||||
def delete_profile(config: ConfigManager, name: str) -> List[str]:
|
||||
"""删除档案,并自动清理其他档案 fallback 中的悬空引用。"""
|
||||
models = get_models(config)
|
||||
if name not in models:
|
||||
return [f"❌ 档案 '{name}' 不存在"]
|
||||
|
||||
del models[name]
|
||||
|
||||
# 清理其他档案对该档案的 fallback 引用(悬空引用会导致启动期报错)
|
||||
cleaned = []
|
||||
for other, profile in models.items():
|
||||
if not isinstance(profile, dict):
|
||||
continue
|
||||
fb = profile.get("fallback")
|
||||
# M4(审计 2026-08-20):fallback 为字符串时 `name in fb` 退化为子串
|
||||
# 误判,且下方列表推导会把字符串拆成单字符数组写回配置
|
||||
if not isinstance(fb, list):
|
||||
continue
|
||||
if name in fb:
|
||||
profile["fallback"] = [f for f in fb if f != name]
|
||||
cleaned.append(other)
|
||||
config._config["models"] = models
|
||||
|
||||
# 若删除的是默认档案,清除默认标记(回退旧 api 节语义)
|
||||
if config.get("api.default_model") == name:
|
||||
_api_section(config).pop("default_model", None)
|
||||
|
||||
saved = config._save_config()
|
||||
msgs = [f"✅ 接入点 '{name}' 已删除"]
|
||||
if cleaned:
|
||||
msgs.append(f"ℹ️ 已同时清理 {len(cleaned)} 个档案的失败转移引用: {', '.join(cleaned)}")
|
||||
if not saved:
|
||||
msgs.append("⚠️ 保存失败——请检查配置文件目录权限")
|
||||
return msgs
|
||||
|
||||
|
||||
def test_profile(config: ConfigManager, name: str, logger=None) -> str:
|
||||
"""连通性验证:向该档案端点发一次最小请求。"""
|
||||
models = get_models(config)
|
||||
if name not in models:
|
||||
return f"❌ 档案 '{name}' 不存在"
|
||||
try:
|
||||
from llm_client import create_llm_client
|
||||
client = create_llm_client(config, api_key=None, logger=logger, model=name)
|
||||
result = client.chat_completion(
|
||||
[{"role": "user", "content": "ping"}], max_tokens=5)
|
||||
if result is None or not str(result).strip():
|
||||
return f"❌ 档案 '{name}' 验证失败:端点返回空响应(多为 4xx 认证/参数错误)"
|
||||
return f"✅ 档案 '{name}' 连通正常:{str(result).strip()[:60]}"
|
||||
except Exception as e:
|
||||
# J4(审计 2026-08-19):异常消息可能含 response.text 片段,网关/WAF
|
||||
# 错误页会回显 Authorization 头——脱敏 Bearer 令牌后再回显终端
|
||||
return f"❌ 档案 '{name}' 验证失败:{_sanitize_log_msg(str(e))}"
|
||||
|
||||
|
||||
# ==================== 交互式菜单 ====================
|
||||
|
||||
def _ask(prompt: str, default: str = "") -> str:
|
||||
if default:
|
||||
val = input(f"{prompt}(回车默认: {default}): ").strip()
|
||||
return val if val else default
|
||||
return input(f"{prompt}: ").strip()
|
||||
|
||||
|
||||
def _ask_yes_no(prompt: str, default: bool = False) -> bool:
|
||||
hint = "y/n" if not default else "Y/n"
|
||||
while True:
|
||||
val = input(f"{prompt} [{hint}] ").strip().lower()
|
||||
if not val:
|
||||
return default
|
||||
if val in ("y", "yes", "是"):
|
||||
return True
|
||||
if val in ("n", "no", "否"):
|
||||
return False
|
||||
print(" 请输入 y 或 n")
|
||||
|
||||
|
||||
def _choose(prompt: str, options: List[str], multi: bool = False) -> List[str]:
|
||||
"""编号选择(单选或多选)。多选用逗号分隔,直接回车 = 不选。"""
|
||||
for i, opt in enumerate(options, 1):
|
||||
print(f" {i}. {opt}")
|
||||
if multi:
|
||||
while True:
|
||||
raw = input(f"{prompt}(多选,逗号分隔,直接回车不选): ").strip()
|
||||
if not raw:
|
||||
return []
|
||||
try:
|
||||
idxs = [int(x) for x in raw.replace(",", ",").split(",") if x.strip()]
|
||||
except ValueError:
|
||||
print(" 请输入数字编号,逗号分隔")
|
||||
continue
|
||||
if all(1 <= i <= len(options) for i in idxs):
|
||||
return [options[i - 1] for i in idxs]
|
||||
print(f" 编号超出范围(1~{len(options)})")
|
||||
else:
|
||||
while True:
|
||||
raw = input(f"{prompt}(1~{len(options)},直接回车不选): ").strip()
|
||||
if not raw:
|
||||
return []
|
||||
try:
|
||||
idx = int(raw)
|
||||
except ValueError:
|
||||
print(" 请输入数字编号")
|
||||
continue
|
||||
if 1 <= idx <= len(options):
|
||||
return [options[idx - 1]]
|
||||
print(f" 编号超出范围(1~{len(options)})")
|
||||
|
||||
|
||||
# T-L17(审计 2026-08-21):按 api_url 域名推荐 Key 环境变量名——此前默认
|
||||
# DEEPSEEK_API_KEY,一路回车会把 DeepSeek 凭据发给其他厂商(moonshot/
|
||||
# aliyuncs 等),正是 resolve_api_key 文档声称要杜绝的跨厂商凭据外发。
|
||||
# 未知域名留空强制输入
|
||||
_ENV_NAME_BY_DOMAIN = {
|
||||
"deepseek": "DEEPSEEK_API_KEY",
|
||||
"moonshot": "MOONSHOT_API_KEY",
|
||||
"aliyuncs": "DASHSCOPE_API_KEY",
|
||||
"bigmodel": "ZHIPU_API_KEY",
|
||||
"ark.cn": "ARK_API_KEY",
|
||||
"openrouter": "OPENROUTER_API_KEY",
|
||||
"openai.com": "OPENAI_API_KEY",
|
||||
"siliconflow": "SILICONFLOW_API_KEY",
|
||||
"volces": "ARK_API_KEY",
|
||||
}
|
||||
|
||||
|
||||
def _suggest_env_name(api_url: str) -> str:
|
||||
url = (api_url or "").lower()
|
||||
for kw, env in _ENV_NAME_BY_DOMAIN.items():
|
||||
if kw in url:
|
||||
return env
|
||||
return ""
|
||||
|
||||
|
||||
def _sanitize_display(s) -> str:
|
||||
"""T-L18(审计 2026-08-21):手写 config.json 可嵌 ANSI/控制转义
|
||||
(\x1b]0;pwned\x07 等),展示层原样 print 每次列表即触发;写入侧有
|
||||
isprintable 校验,读出侧同样过滤控制字符。"""
|
||||
return "".join(ch for ch in str(s or "") if ch.isprintable())
|
||||
|
||||
|
||||
def _print_profiles(profiles: List[Dict]):
|
||||
if not profiles:
|
||||
print(" (当前没有任何接入点。旧版 api 节仍可用,或选 2 添加)")
|
||||
return
|
||||
for p in profiles:
|
||||
mark = " ⭐默认" if p["is_default"] else ""
|
||||
key_desc = f"环境变量 {_sanitize_display(p['api_key_env'])}" if p["api_key_env"] else "免密"
|
||||
fb = (f",失败转移 → "
|
||||
f"{', '.join(_sanitize_display(f) for f in p['fallback'])}") if p["fallback"] else ""
|
||||
print(f" {_sanitize_display(p['name'])}{mark} [{_sanitize_display(p['display_name'])}]")
|
||||
print(f" 地址: {_sanitize_display(p['api_url'])} 模型: {_sanitize_display(p['model'])}")
|
||||
print(f" Key: {key_desc} 温度: {p['temperature'] or '默认'} 超时: {p['timeout'] or '默认'}{fb}")
|
||||
|
||||
|
||||
def _add_flow(config: ConfigManager):
|
||||
print("\n── 添加接入点 ──")
|
||||
print("常见厂商端点参考:")
|
||||
for k, v in PROVIDER_HINTS.items():
|
||||
print(f" {k}: {v}")
|
||||
name = _ask("档案名(英文,如 deepseek/kimi/local)")
|
||||
if not name:
|
||||
print("❌ 已取消(档案名不能为空)")
|
||||
return
|
||||
if not PROFILE_NAME_RE.match(name):
|
||||
print(f"❌ 档案名只能包含字母/数字/下划线/短横线")
|
||||
return
|
||||
if name in get_models(config):
|
||||
print(f"❌ 档案名 '{name}' 已存在")
|
||||
return
|
||||
api_url = _ask("API 地址")
|
||||
model = _ask("模型名(如 deepseek-chat / qwen-plus)")
|
||||
|
||||
print("Key 方式:")
|
||||
print(" 1. 环境变量名(推荐,如 DEEPSEEK_API_KEY——Key 不落盘,重启不丢)")
|
||||
print(" 2. 免密本地端点(如 Ollama,不需要 Key)")
|
||||
key_mode = _ask("选择 [1/2]", "1")
|
||||
api_key_env = ""
|
||||
if key_mode == "1":
|
||||
suggested = _suggest_env_name(api_url)
|
||||
prompt = "环境变量名" + (f"(回车用 {suggested})" if suggested else "")
|
||||
api_key_env = _ask(prompt, suggested)
|
||||
if not api_key_env:
|
||||
print("❌ 已取消(环境变量名不能为空;免密端点请选方式 2)")
|
||||
return
|
||||
elif key_mode != "2":
|
||||
print("❌ 已取消(Key 方式无效)")
|
||||
return
|
||||
|
||||
display_name = _ask("显示名(可选,回车用档案名)")
|
||||
temperature_raw = _ask("温度 0.0~2.0(可选,越低越确定)")
|
||||
timeout_raw = _ask("超时秒数(可选)")
|
||||
reasoning = _ask_yes_no("启用思考模式?", default=False)
|
||||
|
||||
profiles = list_profiles(config)
|
||||
fallback = []
|
||||
if profiles:
|
||||
print("失败转移:主接入点故障时按顺序尝试的备用接入点(可选)")
|
||||
fallback = _choose("选择备用接入点", [p["name"] for p in profiles], multi=True)
|
||||
|
||||
set_default = _ask_yes_no("设为默认接入点?", default=(len(profiles) == 0))
|
||||
|
||||
print("\n── 确认 ──")
|
||||
print(f" 档案名: {name}")
|
||||
print(f" 地址: {api_url} 模型: {model}")
|
||||
print(f" Key: {'环境变量 ' + api_key_env if api_key_env else '免密(不携带)'}")
|
||||
if display_name:
|
||||
print(f" 显示名: {display_name}")
|
||||
if temperature_raw:
|
||||
print(f" 温度: {temperature_raw}")
|
||||
if timeout_raw:
|
||||
print(f" 超时: {timeout_raw}s")
|
||||
if reasoning:
|
||||
print(" 思考模式: 开启")
|
||||
if fallback:
|
||||
print(f" 失败转移: {', '.join(fallback)}")
|
||||
if not _ask_yes_no("确认保存?", default=True):
|
||||
print("已取消")
|
||||
return
|
||||
|
||||
# J8(审计 2026-08-19):原始字符串直传——此前交互层先 float()/int(),
|
||||
# 非数字输入直接 traceback 退出,add_profile 内部的容错成为死代码
|
||||
msgs = add_profile(config, name, api_url, model,
|
||||
api_key_env=api_key_env, display_name=display_name,
|
||||
temperature=temperature_raw or None,
|
||||
timeout=timeout_raw or None,
|
||||
reasoning_enabled=reasoning, fallback=fallback,
|
||||
set_default=set_default)
|
||||
for m in msgs:
|
||||
print(m)
|
||||
|
||||
|
||||
def _delete_flow(config: ConfigManager):
|
||||
profiles = list_profiles(config)
|
||||
if not profiles:
|
||||
print("\n(当前没有任何接入点)")
|
||||
return
|
||||
print("\n── 删除接入点 ──")
|
||||
picked = _choose("选择要删除的接入点", [p["name"] for p in profiles])
|
||||
if not picked:
|
||||
print("已取消")
|
||||
return
|
||||
name = picked[0]
|
||||
if not _ask_yes_no(f"确认删除 '{name}'?(其他档案的失败转移引用会自动清理)", default=False):
|
||||
print("已取消")
|
||||
return
|
||||
for m in delete_profile(config, name):
|
||||
print(m)
|
||||
|
||||
|
||||
def _test_flow(config: ConfigManager):
|
||||
profiles = list_profiles(config)
|
||||
if not profiles:
|
||||
print("\n(当前没有任何接入点)")
|
||||
return
|
||||
print("\n── 连通性验证 ──")
|
||||
picked = _choose("选择要验证的接入点", [p["name"] for p in profiles])
|
||||
if not picked:
|
||||
print("已取消")
|
||||
return
|
||||
print(" 正在请求端点(最长 20 秒)...")
|
||||
print(test_profile(config, picked[0]))
|
||||
|
||||
|
||||
def run_interactive(config: ConfigManager = None):
|
||||
"""交互式主循环。config 可注入(测试用)。"""
|
||||
config = config or ConfigManager()
|
||||
config_path = getattr(config, "_config_path", None) or "(默认位置)"
|
||||
print("=" * 56)
|
||||
print(" aibox 大模型接入点管理工具")
|
||||
print(f" 配置文件: {config_path}")
|
||||
print(" 说明:API Key 只存环境变量名,不落盘;免密端点适合 Ollama 等本地服务")
|
||||
print("=" * 56)
|
||||
while True:
|
||||
# M2(审计 2026-08-20):input 未捕获 EOFError/KeyboardInterrupt——
|
||||
# 管道耗尽/Ctrl+C 此前直接 traceback;整循环包 try 同时覆盖子流程
|
||||
# (_add_flow/_delete_flow/_test_flow)内的 input 调用
|
||||
try:
|
||||
print("\n请选择操作:")
|
||||
print(" 1. 📋 列出接入点")
|
||||
print(" 2. ➕ 添加接入点")
|
||||
print(" 3. 🗑️ 删除接入点")
|
||||
print(" 4. 🔍 验证接入点连通性")
|
||||
print(" 0. 退出")
|
||||
choice = input("> ").strip()
|
||||
if choice == "1":
|
||||
_print_profiles(list_profiles(config))
|
||||
elif choice == "2":
|
||||
_add_flow(config)
|
||||
elif choice == "3":
|
||||
_delete_flow(config)
|
||||
elif choice == "4":
|
||||
_test_flow(config)
|
||||
elif choice in ("0", "q", "exit", "退出"):
|
||||
print("再见")
|
||||
return
|
||||
else:
|
||||
print("无效选择,请输入 0~4")
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
print("\n再见")
|
||||
return
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="aibox 大模型接入点管理工具")
|
||||
parser.add_argument("--list", action="store_true", help="非交互:列出当前接入点")
|
||||
args = parser.parse_args()
|
||||
|
||||
config = ConfigManager()
|
||||
if args.list:
|
||||
_print_profiles(list_profiles(config))
|
||||
return
|
||||
run_interactive(config)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+25
-26
@@ -59,14 +59,14 @@ class AESCli:
|
||||
return False
|
||||
|
||||
def generate_random_key(self, key_size=16):
|
||||
"""生成随机密钥"""
|
||||
"""生成随机密钥(S8:密钥内容不打印,避免进入会话日志;经 --save-key 持久化)"""
|
||||
if key_size not in [16, 24, 32]:
|
||||
print(f"错误: 密钥大小必须为16、24或32,当前为{key_size}")
|
||||
return None
|
||||
|
||||
self.key = get_random_bytes(key_size)
|
||||
key_hex = self.key.hex()
|
||||
print(f"已生成 {key_size} 字节随机密钥: {key_hex}")
|
||||
print(f"已生成 {key_size} 字节随机密钥(内容不显示,请用 --save-key 保存到文件)")
|
||||
return key_hex
|
||||
|
||||
def generate_random_iv(self):
|
||||
@@ -330,6 +330,10 @@ class AESCli:
|
||||
|
||||
with open(filepath, 'w', encoding='utf-8') as f:
|
||||
json.dump(data, f, indent=2, ensure_ascii=False)
|
||||
try:
|
||||
os.chmod(filepath, 0o600) # S8:密钥文件仅属主可读写
|
||||
except OSError:
|
||||
pass # 不支持 POSIX 权限的平台(如 Windows)跳过
|
||||
|
||||
print(f"密钥和IV已保存到: {filepath}")
|
||||
return True
|
||||
@@ -395,38 +399,33 @@ def main():
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
使用示例:
|
||||
# 生成随机密钥和IV
|
||||
aes-crypto --gen-key
|
||||
# 生成随机密钥和IV(密钥不打印,经 --save-key 保存到文件)
|
||||
aes-crypto --gen-key --gen-iv --save-key keys.json
|
||||
aes-crypto --gen-iv
|
||||
|
||||
# 文本加解密
|
||||
aes-crypto --key <密钥hex> --iv <IV/Nonce hex> --encrypt-text "Hello World"
|
||||
aes-crypto --key <密钥hex> --iv <IV/Nonce hex> --decrypt-text "<密文base64>"
|
||||
# 文本加解密(密钥经文件提供,不走命令行)
|
||||
aes-crypto --key-file key.txt --iv <IV/Nonce hex> --encrypt-text "Hello World"
|
||||
aes-crypto --key-file key.txt --iv <IV/Nonce hex> --decrypt-text "<密文base64>"
|
||||
|
||||
# 文件加解密
|
||||
aes-crypto --key <密钥hex> --iv <IV/Nonce hex> --encrypt-file input.txt
|
||||
aes-crypto --key <密钥hex> --iv <IV/Nonce hex> --decrypt-file input.txt.encrypted
|
||||
aes-crypto --load-key keys.json --encrypt-file input.txt
|
||||
aes-crypto --load-key keys.json --decrypt-file input.txt.encrypted
|
||||
|
||||
# 批量加解密
|
||||
aes-crypto --key <密钥hex> --iv <IV/Nonce hex> --batch-encrypt file1.txt file2.txt
|
||||
aes-crypto --key <密钥hex> --iv <IV/Nonce hex> --batch-decrypt *.encrypted
|
||||
|
||||
# 保存和加载密钥
|
||||
aes-crypto --key <密钥hex> --iv <IV/Nonce hex> --save-key keys.json
|
||||
aes-crypto --load-key keys.json --encrypt-file input.txt
|
||||
aes-crypto --load-key keys.json --batch-encrypt file1.txt file2.txt
|
||||
aes-crypto --load-key keys.json --batch-decrypt *.encrypted
|
||||
|
||||
# 使用GCM模式
|
||||
aes-crypto --mode GCM --gen-key --key-size 32
|
||||
aes-crypto --mode GCM --key <密钥hex> --iv <Nonce hex> --encrypt-file input.txt
|
||||
aes-crypto --mode GCM --gen-key --key-size 32 --save-key keys.json
|
||||
aes-crypto --mode GCM --load-key keys.json --encrypt-file input.txt
|
||||
|
||||
# 从文件读取密钥
|
||||
aes-crypto --key-file key.txt --iv <IV/Nonce hex> --encrypt-file input.txt
|
||||
|
||||
参数说明:
|
||||
--mode CBC|GCM 加密模式 (默认: CBC)
|
||||
--key KEY 十六进制密钥 (16/24/32字节)
|
||||
--iv IV 十六进制IV/Nonce (CBC:16字节, GCM:12字节)
|
||||
--key-file FILE 从文件读取密钥
|
||||
--key-file FILE 从文件读取密钥(密钥仅经文件提供,避免进入进程列表与会话日志)
|
||||
--load-key FILE 从JSON文件加载密钥和IV
|
||||
--save-key FILE 保存密钥和IV到JSON文件
|
||||
--gen-key 生成随机密钥
|
||||
@@ -448,8 +447,7 @@ def main():
|
||||
parser.add_argument('--mode', choices=['CBC', 'GCM'], default='CBC',
|
||||
help='加密模式 (默认: CBC)')
|
||||
|
||||
# 密钥和IV设置
|
||||
parser.add_argument('--key', help='十六进制密钥 (16/24/32字节)')
|
||||
# 密钥和IV设置(S7:密钥仅经 --key-file/--load-key 文件途径提供,不进命令行)
|
||||
parser.add_argument('--iv', help='十六进制IV/Nonce (CBC:16字节, GCM:12字节)')
|
||||
parser.add_argument('--key-file', help='从文件读取密钥')
|
||||
parser.add_argument('--load-key', help='从JSON文件加载密钥和IV')
|
||||
@@ -491,27 +489,28 @@ def main():
|
||||
sys.exit(1)
|
||||
|
||||
# 处理密钥文件
|
||||
key_input = None
|
||||
if args.key_file:
|
||||
try:
|
||||
with open(args.key_file, 'r') as f:
|
||||
args.key = f.read().strip()
|
||||
key_input = f.read().strip()
|
||||
except Exception as e:
|
||||
print(f"错误: 读取密钥文件失败: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
# 设置密钥和IV
|
||||
if args.key:
|
||||
if key_input:
|
||||
if args.iv is None and not aes.iv:
|
||||
print("错误: 需要同时提供密钥和IV/Nonce")
|
||||
sys.exit(1)
|
||||
|
||||
if args.iv:
|
||||
if not aes.set_key_iv(args.key, args.iv):
|
||||
if not aes.set_key_iv(key_input, args.iv):
|
||||
sys.exit(1)
|
||||
elif aes.iv:
|
||||
# 从加载的密钥文件中获取IV
|
||||
try:
|
||||
aes.key = bytes.fromhex(args.key)
|
||||
aes.key = bytes.fromhex(key_input)
|
||||
if len(aes.key) not in [16, 24, 32]:
|
||||
print(f"错误: 密钥长度必须为16、24或32字节,当前为{len(aes.key)}字节")
|
||||
sys.exit(1)
|
||||
@@ -548,7 +547,7 @@ def main():
|
||||
if args.encrypt_text or args.decrypt_text or \
|
||||
args.encrypt_file or args.decrypt_file or \
|
||||
args.batch_encrypt or args.batch_decrypt:
|
||||
print("错误: 请先设置密钥和IV/Nonce (使用 --key 和 --iv 或 --load-key)")
|
||||
print("错误: 请先设置密钥和IV/Nonce (使用 --key-file 和 --iv 或 --load-key)")
|
||||
sys.exit(1)
|
||||
|
||||
try:
|
||||
|
||||
@@ -18,6 +18,16 @@ def ensure_camera_dir():
|
||||
return CAMERA_DIR
|
||||
|
||||
|
||||
def validate_filename(filename):
|
||||
"""S9(审计 2026-08-20):--filename 仅允许纯文件名——
|
||||
拒绝路径分隔符/父目录引用/绝对路径,防止经 os.path.join 逃逸 CAMERA_DIR"""
|
||||
if not filename:
|
||||
return "文件名不能为空"
|
||||
if filename in (".", "..") or "/" in filename or "\\" in filename or os.path.isabs(filename):
|
||||
return f"非法文件名(不允许路径分隔符或父目录引用): {filename}"
|
||||
return None
|
||||
|
||||
|
||||
def list_cameras(max_test=10):
|
||||
"""列出可用的摄像头设备"""
|
||||
available_cameras = []
|
||||
@@ -341,6 +351,10 @@ def main():
|
||||
ensure_camera_dir()
|
||||
|
||||
if args.filename:
|
||||
filename_error = validate_filename(args.filename)
|
||||
if filename_error:
|
||||
print(f"❌ {filename_error}")
|
||||
sys.exit(1)
|
||||
save_path = os.path.join(CAMERA_DIR, args.filename)
|
||||
else:
|
||||
timestamp = time.strftime("%Y%m%d_%H%M%S")
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
"""人脸检测和识别工具 - 检测图像中的人脸并保存裁剪的人脸图像。"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import argparse
|
||||
import cv2
|
||||
@@ -15,6 +16,16 @@ except ImportError:
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
# S11(审计 2026-08-20):person_name 白名单——仅允许文字/数字/._-,
|
||||
# 拒绝路径分隔符与纯点名("."/".."),防止经 f"{name}.jpg" 拼接逃逸 known_faces 目录。
|
||||
# \w 在 py3 默认 UNICODE,兼容中文人名(如 "张三")。
|
||||
_PERSON_NAME_RE = re.compile(r"^(?!\.+$)[\w.-]{1,64}$")
|
||||
|
||||
|
||||
def _valid_person_name(name):
|
||||
return bool(name) and bool(_PERSON_NAME_RE.match(name))
|
||||
|
||||
|
||||
class FaceDetector:
|
||||
"""人脸检测和识别工具"""
|
||||
|
||||
@@ -198,6 +209,8 @@ class FaceDetector:
|
||||
def add_known_face(self, image_path: str, person_name: str) -> str:
|
||||
"""添加已知人脸到数据库"""
|
||||
try:
|
||||
if not _valid_person_name(person_name):
|
||||
return f"❌ 非法人名标识: {person_name}(仅允许文字/数字/._-,≤64字符)"
|
||||
if not os.path.exists(image_path):
|
||||
return f"❌ 文件不存在: {image_path}"
|
||||
|
||||
@@ -237,6 +250,8 @@ class FaceDetector:
|
||||
def remove_known_face(self, person_name: str) -> str:
|
||||
"""从已知人脸库中删除人脸"""
|
||||
try:
|
||||
if not _valid_person_name(person_name):
|
||||
return f"❌ 非法人名标识: {person_name}(仅允许文字/数字/._-,≤64字符)"
|
||||
filepath = os.path.join(self.known_faces_dir, f"{person_name}.jpg")
|
||||
|
||||
if not os.path.exists(filepath):
|
||||
@@ -260,6 +275,8 @@ class FaceDetector:
|
||||
def get_face_info(self, person_name: str) -> str:
|
||||
"""获取已知人脸信息"""
|
||||
try:
|
||||
if not _valid_person_name(person_name):
|
||||
return f"❌ 非法人名标识: {person_name}(仅允许文字/数字/._-,≤64字符)"
|
||||
filepath = os.path.join(self.known_faces_dir, f"{person_name}.jpg")
|
||||
|
||||
if not os.path.exists(filepath):
|
||||
|
||||
+41
-10
@@ -3,6 +3,7 @@
|
||||
|
||||
import sys
|
||||
import os
|
||||
import re
|
||||
import argparse
|
||||
import requests
|
||||
from datetime import datetime
|
||||
@@ -35,6 +36,13 @@ MAX_CONTENT_PREVIEW = 1000
|
||||
MAX_LIST_PREVIEW = 20
|
||||
MAX_STRING_LENGTH = 200
|
||||
|
||||
# S3:分支/ref 白名单——首字符必须为字母或数字,防止以 '-' 开头的值被 git 解析为选项(选项注入)
|
||||
_REF_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._/@+-]{0,127}$")
|
||||
|
||||
|
||||
def _valid_ref(name: str) -> bool:
|
||||
return bool(name) and bool(_REF_RE.match(name))
|
||||
|
||||
|
||||
# ==================== GitHub工具类 ====================
|
||||
|
||||
@@ -79,19 +87,30 @@ class GitHubTool:
|
||||
return date_str
|
||||
|
||||
def _run_git(self, repo_path: str, cmd: List[str]) -> Dict:
|
||||
env = {**os.environ, "GIT_TERMINAL_PROMPT": "0"}
|
||||
if GITHUB_TOKEN:
|
||||
# 凭据仅经环境变量传递(git>=2.31 的 GIT_CONFIG_* 机制),
|
||||
# 不进命令行 argv / 远端 URL / .git/config(S1,对齐 J12 条例)
|
||||
env["GIT_CONFIG_COUNT"] = "1"
|
||||
env["GIT_CONFIG_KEY_0"] = "http.extraheader"
|
||||
env["GIT_CONFIG_VALUE_0"] = f"Authorization: token {GITHUB_TOKEN}"
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["git"] + cmd, cwd=repo_path, capture_output=True,
|
||||
text=True, timeout=60, env={**os.environ, "GIT_TERMINAL_PROMPT": "0"}
|
||||
text=True, timeout=60, env=env
|
||||
)
|
||||
return {"success": result.returncode == 0, "stdout": result.stdout.strip(), "stderr": result.stderr.strip()}
|
||||
return {"success": result.returncode == 0,
|
||||
"stdout": self._scrub_token(result.stdout.strip()),
|
||||
"stderr": self._scrub_token(result.stderr.strip())}
|
||||
except Exception as e:
|
||||
return {"success": False, "error": str(e)}
|
||||
return {"success": False, "error": self._scrub_token(str(e))}
|
||||
|
||||
def _auth_url(self, repo_url: str) -> str:
|
||||
if GITHUB_TOKEN and repo_url.startswith("https://"):
|
||||
return repo_url.replace("https://", f"https://{GITHUB_TOKEN}@")
|
||||
return repo_url
|
||||
@staticmethod
|
||||
def _scrub_token(text: str) -> str:
|
||||
"""输出脱敏:任何路径回显的 token 一律替换(S1)"""
|
||||
if GITHUB_TOKEN and text:
|
||||
return text.replace(GITHUB_TOKEN, "***")
|
||||
return text
|
||||
|
||||
# ==================== 搜索功能 ====================
|
||||
|
||||
@@ -154,7 +173,9 @@ class GitHubTool:
|
||||
if os.path.exists(target_dir):
|
||||
return f"⚠️ 仓库已存在: {target_dir}\n💡 使用 pull 命令拉取更新"
|
||||
|
||||
clone_url = self._auth_url(f"https://github.com/{repo}.git")
|
||||
if branch and not _valid_ref(branch):
|
||||
return f"❌ 非法分支名: {branch}"
|
||||
clone_url = f"https://github.com/{repo}.git"
|
||||
cmd = ["clone"]
|
||||
if depth > 0:
|
||||
cmd.extend(["--depth", str(depth)])
|
||||
@@ -175,6 +196,8 @@ class GitHubTool:
|
||||
br = self._run_git(repo_path, ["branch", "--show-current"])
|
||||
branch = br["stdout"] if br["success"] else None
|
||||
|
||||
if branch and not _valid_ref(branch):
|
||||
return f"❌ 非法分支名: {branch}"
|
||||
result = self._run_git(repo_path, ["pull", "origin", branch] if branch else ["pull"])
|
||||
if result["success"]:
|
||||
return f"✅ 拉取成功\n {result['stdout'] or '已是最新'}"
|
||||
@@ -188,6 +211,8 @@ class GitHubTool:
|
||||
br = self._run_git(repo_path, ["branch", "--show-current"])
|
||||
branch = br["stdout"] if br["success"] else None
|
||||
|
||||
if branch and not _valid_ref(branch):
|
||||
return f"❌ 非法分支名: {branch}"
|
||||
cmd = ["push"]
|
||||
if force:
|
||||
cmd.append("--force")
|
||||
@@ -273,9 +298,15 @@ class GitHubTool:
|
||||
return "⚠️ 需要确认删除,请添加 --confirm 参数"
|
||||
if not os.path.exists(repo_path):
|
||||
return f"❌ 路径不存在: {repo_path}"
|
||||
# S2:仅允许删除 WORKSPACE 内的 git 仓库,防止任意目录递归删除
|
||||
workspace = os.path.realpath(WORKSPACE_DIR)
|
||||
real_path = os.path.realpath(repo_path)
|
||||
if not real_path.startswith(workspace + os.sep) \
|
||||
or not os.path.isdir(os.path.join(real_path, ".git")):
|
||||
return f"❌ 仅允许删除工作区({workspace})内的 git 仓库目录"
|
||||
try:
|
||||
shutil.rmtree(repo_path)
|
||||
return f"✅ 已删除: {repo_path}"
|
||||
shutil.rmtree(real_path)
|
||||
return f"✅ 已删除: {real_path}"
|
||||
except Exception as e:
|
||||
return f"❌ 删除失败: {e}"
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
一个基于python-nmap的局域网Ping扫描工具
|
||||
"""
|
||||
import argparse
|
||||
import shlex
|
||||
import sys
|
||||
import socket
|
||||
|
||||
@@ -25,7 +26,9 @@ def scan_network(network="192.168.1.0/24", timeout=3, verbose=False):
|
||||
print("-" * 50)
|
||||
|
||||
try:
|
||||
result = nm.scan(hosts=network, arguments=arguments)
|
||||
# J12(审计 2026-08-19):python-nmap 内部把 hosts 拼进 shell 命令
|
||||
# 执行——LLM/AI 可控的 network 参数需 shlex.quote,防参数/命令注入
|
||||
result = nm.scan(hosts=shlex.quote(network), arguments=arguments)
|
||||
except Exception as e:
|
||||
print(f"❌ 扫描失败: {e}")
|
||||
print("💡 提示:请确保已安装 nmap (sudo apt install nmap) 并有足够权限")
|
||||
|
||||
@@ -29,6 +29,14 @@ def get_timestamp_filename(prefix="screenshot", extension="png"):
|
||||
return f"{prefix}_{timestamp}.{extension}"
|
||||
|
||||
|
||||
def validate_filename(filename):
|
||||
"""S10(审计 2026-08-20):--filename 仅允许纯文件名——
|
||||
拒绝路径分隔符/父目录引用/绝对路径,防止经 os.path.join 逃逸保存目录"""
|
||||
if filename in (".", "..") or "/" in filename or "\\" in filename or os.path.isabs(filename):
|
||||
return f"非法文件名(不允许路径分隔符或父目录引用): {filename}"
|
||||
return None
|
||||
|
||||
|
||||
def get_system():
|
||||
"""获取当前操作系统"""
|
||||
return platform.system()
|
||||
@@ -41,6 +49,10 @@ def take_screenshot_windows(output_path, delay=0):
|
||||
if delay > 0:
|
||||
time.sleep(delay)
|
||||
|
||||
# J12(审计 2026-08-19):output_path 经 PowerShell 单引号转义——
|
||||
# 路径中的单引号翻倍(PS 字符串转义),AI/用户可控文件名无法闭合
|
||||
# 字符串注入脚本(此前 `${output_path}` 直拼,`");calc;#` 即注入)
|
||||
ps_path = str(output_path).replace("'", "''")
|
||||
try:
|
||||
ps_script = f'''
|
||||
Add-Type -AssemblyName System.Windows.Forms
|
||||
@@ -53,11 +65,11 @@ $bitmap = New-Object System.Drawing.Bitmap $bounds.Width, $bounds.Height
|
||||
$graphics = [System.Drawing.Graphics]::FromImage($bitmap)
|
||||
$graphics.CopyFromScreen($bounds.X, $bounds.Y, 0, 0, $bounds.Size)
|
||||
|
||||
$bitmap.Save("{output_path}", [System.Drawing.Imaging.ImageFormat]::Png)
|
||||
$bitmap.Save('{ps_path}', [System.Drawing.Imaging.ImageFormat]::Png)
|
||||
$graphics.Dispose()
|
||||
$bitmap.Dispose()
|
||||
|
||||
Write-Output "截图已保存: {output_path}"
|
||||
Write-Output "截图已保存: {ps_path}"
|
||||
'''
|
||||
|
||||
result = subprocess.run(
|
||||
@@ -150,6 +162,9 @@ def take_screenshot(output_dir=None, filename=None, delay=0):
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
if filename:
|
||||
filename_error = validate_filename(filename)
|
||||
if filename_error:
|
||||
return False, filename_error, None
|
||||
filepath = os.path.join(save_dir, filename)
|
||||
else:
|
||||
filepath = os.path.join(save_dir, get_timestamp_filename())
|
||||
|
||||
@@ -27,9 +27,12 @@ def build_parser():
|
||||
--host SSH服务器主机名或IP地址 (必填)
|
||||
--port SSH服务器端口号,默认 22
|
||||
-u, --username SSH登录用户名 (必填)
|
||||
-p, --password SSH登录密码 (如不提供则交互式输入)
|
||||
--key SSH私钥文件路径 (用于密钥认证)
|
||||
|
||||
密码认证:
|
||||
密码经环境变量 AIBOX_SFTP_PASS 提供;未设置且未用 --key 时
|
||||
交互式输入。不走命令行参数,避免凭据进入进程列表与会话日志(J12 条例)。
|
||||
|
||||
操作参数:
|
||||
--upload LOCAL REMOTE 上传:本地文件/目录 -> 远程服务器路径
|
||||
--download REMOTE LOCAL 下载:远程服务器文件/目录 -> 本地路径
|
||||
@@ -37,19 +40,19 @@ def build_parser():
|
||||
-v, --verbose 显示详细信息
|
||||
|
||||
使用示例:
|
||||
1. 上传文件(密码认证):
|
||||
sftp-transfer --host 192.168.1.100 -u root -p mypass --upload /tmp/test.txt /root/
|
||||
1. 上传文件(密码认证,环境变量提供密码):
|
||||
AIBOX_SFTP_PASS=mypass sftp-transfer --host 192.168.1.100 -u root --upload /tmp/test.txt /root/
|
||||
|
||||
2. 下载文件(密钥认证):
|
||||
sftp-transfer --host 192.168.1.100 -u root --key ~/.ssh/id_rsa \\
|
||||
--download /root/test.txt /tmp/
|
||||
|
||||
3. 递归上传目录:
|
||||
sftp-transfer --host 192.168.1.100 -u user -p pass -r \\
|
||||
AIBOX_SFTP_PASS=pass sftp-transfer --host 192.168.1.100 -u user -r \\
|
||||
--upload ./mydir /home/user/mydir
|
||||
|
||||
4. 下载目录:
|
||||
sftp-transfer --host 192.168.1.100 -u user -p pass -r \\
|
||||
AIBOX_SFTP_PASS=pass sftp-transfer --host 192.168.1.100 -u user -r \\
|
||||
--download /home/user/backup ./backup
|
||||
|
||||
5. 交互式输入密码:
|
||||
@@ -61,7 +64,6 @@ def build_parser():
|
||||
parser.add_argument('--host', required=True, help='SSH服务器主机名或IP地址')
|
||||
parser.add_argument('--port', type=int, default=22, help='SSH服务器端口号,默认22')
|
||||
parser.add_argument('-u', '--username', required=True, help='SSH登录用户名')
|
||||
parser.add_argument('-p', '--password', help='SSH登录密码(如不提供则交互式输入)')
|
||||
parser.add_argument('--key', help='SSH私钥文件路径(用于密钥认证)')
|
||||
|
||||
# 操作参数
|
||||
@@ -81,10 +83,20 @@ def log(verbose, msg):
|
||||
print(f"[SFTP] {msg}")
|
||||
|
||||
|
||||
def safe_member_name(name):
|
||||
"""S6(审计 2026-08-20):目录成员名校验——拒绝路径分隔符与 . / ..,
|
||||
防止恶意 SFTP 服务端构造 listdir 条目(如 ../../etc/cron.d/x)在递归下载时逃逸本地目录"""
|
||||
return name not in ('', '.', '..') and '/' not in name and '\\' not in name
|
||||
|
||||
|
||||
def get_password(args):
|
||||
"""获取密码:优先从命令行参数获取,否则交互式输入"""
|
||||
if args.password:
|
||||
return args.password
|
||||
"""获取密码:环境变量 AIBOX_SFTP_PASS 优先,否则交互式输入(S5:凭据不走命令行)"""
|
||||
env_pass = os.environ.get("AIBOX_SFTP_PASS")
|
||||
if env_pass:
|
||||
return env_pass
|
||||
if args.key:
|
||||
# 密钥认证时不强制索要密码(仅在私钥有口令时经交互输入)
|
||||
return None
|
||||
return getpass.getpass(f"请输入 {args.username}@{args.host} 的密码: ")
|
||||
|
||||
|
||||
@@ -143,6 +155,9 @@ def upload_dir(sftp, local_dir, remote_dir, verbose):
|
||||
"""递归上传目录"""
|
||||
ensure_remote_dir(sftp, remote_dir, verbose)
|
||||
for item in os.listdir(local_dir):
|
||||
if not safe_member_name(item):
|
||||
print(f" 警告: 跳过非法本地条目名: {item!r}", file=sys.stderr)
|
||||
continue
|
||||
local_item = os.path.join(local_dir, item)
|
||||
remote_item = os.path.join(remote_dir, item)
|
||||
if os.path.isdir(local_item):
|
||||
@@ -168,6 +183,9 @@ def download_dir(sftp, remote_dir, local_dir, verbose):
|
||||
log(verbose, f" 创建本地目录: {local_dir}")
|
||||
|
||||
for item in sftp.listdir(remote_dir):
|
||||
if not safe_member_name(item):
|
||||
print(f" 警告: 跳过非法远端条目名: {item!r}", file=sys.stderr)
|
||||
continue
|
||||
remote_item = os.path.join(remote_dir, item).replace('\\', '/')
|
||||
local_item = os.path.join(local_dir, item)
|
||||
try:
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
"""ssh-exec - 通过SSH远程执行单条命令的工具(支持密码认证)"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import paramiko
|
||||
|
||||
@@ -14,24 +15,32 @@ def main():
|
||||
-H, --host SSH主机地址(默认: 127.0.0.1)
|
||||
-P, --port SSH端口号(默认: 22)
|
||||
-u, --user SSH用户名(必填)
|
||||
-p, --password SSH密码(必填)
|
||||
-c, --command 要执行的远程命令(必填)
|
||||
|
||||
密码认证:
|
||||
密码经环境变量 AIBOX_SSH_PASS 提供(必填),不走命令行参数,
|
||||
避免凭据进入进程列表与会话日志(J12 条例)。
|
||||
|
||||
使用示例:
|
||||
%(prog)s -u qing -p 12345678 -c "ls -la"
|
||||
%(prog)s -u root -H 192.168.1.100 -P 22 -p mypass -c "df -h"
|
||||
%(prog)s -u qing -H 127.0.0.1 -p 12345678 -c "uname -a"
|
||||
AIBOX_SSH_PASS=12345678 %(prog)s -u qing -c "ls -la"
|
||||
AIBOX_SSH_PASS=mypass %(prog)s -u root -H 192.168.1.100 -P 22 -c "df -h"
|
||||
AIBOX_SSH_PASS=12345678 %(prog)s -u qing -H 127.0.0.1 -c "uname -a"
|
||||
""",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
parser.add_argument("-H", "--host", default="127.0.0.1", help="SSH主机地址(默认: 127.0.0.1)")
|
||||
parser.add_argument("-P", "--port", type=int, default=22, help="SSH端口号(默认: 22)")
|
||||
parser.add_argument("-u", "--user", required=True, help="SSH用户名(必填)")
|
||||
parser.add_argument("-p", "--password", required=True, help="SSH密码(必填)")
|
||||
parser.add_argument("-c", "--command", required=True, help="要执行的远程命令(必填)")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# 密码仅从环境变量读取(S4:凭据不进命令行)
|
||||
password = os.environ.get("AIBOX_SSH_PASS")
|
||||
if not password:
|
||||
print("错误: 请通过环境变量 AIBOX_SSH_PASS 提供SSH密码", file=sys.stderr)
|
||||
sys.exit(2)
|
||||
|
||||
client = paramiko.SSHClient()
|
||||
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
||||
|
||||
@@ -40,7 +49,7 @@ def main():
|
||||
hostname=args.host,
|
||||
port=args.port,
|
||||
username=args.user,
|
||||
password=args.password,
|
||||
password=password,
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
|
||||
@@ -52,6 +52,13 @@ create_apprun() {
|
||||
local exec_name="$2"
|
||||
local terminal="$3"
|
||||
|
||||
# S13(审计 2026-08-20):exec_name 进入 sed 替换串,含 / & 等字符会注入或破坏
|
||||
# 替换表达式(如 "a/x" 使 s///g 错位、"&" 展开为匹配串)。白名单拦截,拒绝即退出。
|
||||
if [[ ! "$exec_name" =~ ^[A-Za-z0-9._+-]+$ ]]; then
|
||||
echo -e "${RED}错误: 可执行文件名含非法字符(仅允许 A-Za-z0-9._+-): $exec_name${NC}"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
cat > "$app_dir/AppRun" << 'EOF'
|
||||
#!/bin/bash
|
||||
HERE="$(dirname "$(readlink -f "$0")")"
|
||||
|
||||
+26
-1
@@ -23,6 +23,27 @@ def get_api_key():
|
||||
return None
|
||||
|
||||
|
||||
def detect_image_mime(image_path):
|
||||
"""S12(审计 2026-08-20):按魔数识别真实图片类型——
|
||||
拒绝非图片文件(防任意文件读取外泄),并返回正确 MIME(此前硬编码 image/jpeg)"""
|
||||
try:
|
||||
with open(image_path, 'rb') as f:
|
||||
head = f.read(16)
|
||||
except Exception:
|
||||
return None
|
||||
if head.startswith(b'\xff\xd8\xff'):
|
||||
return 'image/jpeg'
|
||||
if head.startswith(b'\x89PNG\r\n\x1a\n'):
|
||||
return 'image/png'
|
||||
if head[:6] in (b'GIF87a', b'GIF89a'):
|
||||
return 'image/gif'
|
||||
if head.startswith(b'RIFF') and head[8:12] == b'WEBP':
|
||||
return 'image/webp'
|
||||
if head.startswith(b'BM'):
|
||||
return 'image/bmp'
|
||||
return None
|
||||
|
||||
|
||||
def encode_image_to_base64(image_path):
|
||||
"""将图片编码为base64"""
|
||||
with open(image_path, 'rb') as f:
|
||||
@@ -38,6 +59,10 @@ def analyze_image(image_path, prompt, stream=False):
|
||||
if not os.path.exists(image_path):
|
||||
return f"❌ 错误:图片文件不存在 - {image_path}"
|
||||
|
||||
mime = detect_image_mime(image_path)
|
||||
if not mime:
|
||||
return f"❌ 错误:文件不是有效图片(仅支持 JPEG/PNG/GIF/WebP/BMP) - {image_path}"
|
||||
|
||||
try:
|
||||
base64_image = encode_image_to_base64(image_path)
|
||||
except Exception as e:
|
||||
@@ -55,7 +80,7 @@ def analyze_image(image_path, prompt, stream=False):
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/jpeg;base64,{base64_image}"
|
||||
"url": f"data:{mime};base64,{base64_image}"
|
||||
}
|
||||
},
|
||||
{
|
||||
|
||||
@@ -24,6 +24,11 @@ from .db import MemoryDatabase, MemoDatabase
|
||||
from .skill_manager import SkillManager
|
||||
from .conversation_history import ConversationHistory
|
||||
from .conversation_compressor import ConversationCompressor
|
||||
from .process_manager import (
|
||||
ProcessManager,
|
||||
get_process_manager,
|
||||
reset_process_manager,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"CommandExecutor",
|
||||
@@ -34,4 +39,7 @@ __all__ = [
|
||||
"SkillManager",
|
||||
"ConversationHistory",
|
||||
"ConversationCompressor",
|
||||
"ProcessManager",
|
||||
"get_process_manager",
|
||||
"reset_process_manager",
|
||||
]
|
||||
@@ -223,49 +223,6 @@ class ConversationCompressor:
|
||||
|
||||
return conversation_turns
|
||||
|
||||
def _is_range_safe(self, messages: List[Dict], start: int, end: int, protected: set) -> bool:
|
||||
"""检查消息范围是否安全 - 简化版,不检查tool完整性"""
|
||||
return True # 不保护tool消息,允许压缩
|
||||
|
||||
def _get_protected_indices(self, messages: List[Dict]) -> set:
|
||||
"""
|
||||
找出所有受保护的消息索引 - 简化版,只保护最近的对话轮次
|
||||
"""
|
||||
protected_indices = set()
|
||||
|
||||
# ========== 1. 保护最近的消息(由配置决定) ==========
|
||||
keep_recent = constants.COMPRESS_KEEP_RECENT
|
||||
if keep_recent > 0:
|
||||
recent_start = max(0, len(messages) - keep_recent)
|
||||
for i in range(recent_start, len(messages)):
|
||||
protected_indices.add(i)
|
||||
|
||||
# ========== 2. 保护前2个对话轮次 ==========
|
||||
conversation_turns = self._build_conversation_turns(messages)
|
||||
|
||||
if len(conversation_turns) >= 2:
|
||||
# 保护第1轮
|
||||
first_turn = conversation_turns[0]
|
||||
protected_indices.add(first_turn["user_index"])
|
||||
if first_turn["assistant_index"] is not None:
|
||||
protected_indices.add(first_turn["assistant_index"])
|
||||
for tool_idx in first_turn["tool_indices"]:
|
||||
protected_indices.add(tool_idx)
|
||||
|
||||
# 保护第2轮
|
||||
second_turn = conversation_turns[1]
|
||||
protected_indices.add(second_turn["user_index"])
|
||||
if second_turn["assistant_index"] is not None:
|
||||
protected_indices.add(second_turn["assistant_index"])
|
||||
for tool_idx in second_turn["tool_indices"]:
|
||||
protected_indices.add(tool_idx)
|
||||
|
||||
# 调试日志
|
||||
if self.logger.level >= 3:
|
||||
self.logger.debug(f"受保护的消息索引: {sorted(protected_indices)}")
|
||||
|
||||
return protected_indices
|
||||
|
||||
def _build_compression_prompt(self, messages: List[Dict]) -> str:
|
||||
"""构建压缩提示词 - 包含工具调用信息"""
|
||||
conversation_lines = []
|
||||
@@ -368,25 +325,57 @@ class ConversationCompressor:
|
||||
return None
|
||||
|
||||
def _validate_message_structure(self, messages: List[Dict]) -> List[Dict]:
|
||||
"""验证并修复消息结构,移除孤立的 tool 消息"""
|
||||
"""验证并修复消息结构(tool 配对双向完整性)。
|
||||
|
||||
OpenAI 协议要求 assistant.tool_calls 与对应的 tool 结果消息成对出现,
|
||||
缺任一侧后续 API 调用直接 400:
|
||||
1. 孤立 tool 消息(tool_call_id 无对应 assistant)→ 移除(既有行为)
|
||||
2. T23(审计 2026-08-21):压缩块边界按条数切,可能落在带 tool_calls
|
||||
的 assistant 消息与其 tool 结果之间——保留侧 assistant 的
|
||||
tool_calls 失去配对 → 剥离失配的 tool_calls 条目;全部失配时删除
|
||||
字段(content 为空则置空串,防产生非法空消息)
|
||||
"""
|
||||
if not messages:
|
||||
return messages
|
||||
|
||||
# 找出所有 assistant 消息的 tool_call_ids
|
||||
assistant_tool_ids = set()
|
||||
# tool 侧实际存在的 id
|
||||
present_tool_ids = {
|
||||
m.get("tool_call_id")
|
||||
for m in messages
|
||||
if m.get("role") == "tool" and m.get("tool_call_id")
|
||||
}
|
||||
|
||||
# 第一遍:assistant 剥离失去配对的 tool_calls
|
||||
first_pass = []
|
||||
kept_assistant_ids = set()
|
||||
for msg in messages:
|
||||
if msg.get("role") == "assistant" and msg.get("tool_calls"):
|
||||
for tc in msg.get("tool_calls", []):
|
||||
tc_id = tc.get("id")
|
||||
if tc_id:
|
||||
assistant_tool_ids.add(tc_id)
|
||||
calls = msg.get("tool_calls", [])
|
||||
kept_calls = [tc for tc in calls
|
||||
if tc.get("id") and tc.get("id") in present_tool_ids]
|
||||
kept_assistant_ids.update(tc.get("id") for tc in kept_calls)
|
||||
if len(kept_calls) == len(calls):
|
||||
first_pass.append(msg) # 全部有配对,原样保留
|
||||
continue
|
||||
new_msg = dict(msg)
|
||||
if kept_calls:
|
||||
new_msg["tool_calls"] = kept_calls
|
||||
print(f" ⚠️ 剥离部分失去配对 tool 结果的 tool_calls(T23)")
|
||||
else:
|
||||
new_msg.pop("tool_calls", None)
|
||||
if not new_msg.get("content"):
|
||||
new_msg["content"] = ""
|
||||
print(f" ⚠️ 剥离失去配对 tool 结果的 assistant.tool_calls(T23)")
|
||||
first_pass.append(new_msg)
|
||||
else:
|
||||
first_pass.append(msg)
|
||||
|
||||
# 过滤掉没有对应 assistant 的 tool 消息
|
||||
# 第二遍:移除孤立 tool 消息(其 id 不在任何保留的 assistant tool_calls 中)
|
||||
filtered = []
|
||||
for msg in messages:
|
||||
for msg in first_pass:
|
||||
if msg.get("role") == "tool":
|
||||
tool_call_id = msg.get("tool_call_id")
|
||||
if tool_call_id and tool_call_id not in assistant_tool_ids:
|
||||
if tool_call_id and tool_call_id not in kept_assistant_ids:
|
||||
print(f" ⚠️ 移除孤立的 tool 消息: {tool_call_id}")
|
||||
continue
|
||||
filtered.append(msg)
|
||||
@@ -514,7 +503,8 @@ class ConversationCompressor:
|
||||
final_tokens = self.estimate_tokens([m for m in final_messages if m.get("role") != "system"])
|
||||
|
||||
print(f"\n✅ {i18n.get('conversation_compressed', original_count, final_count)}")
|
||||
print(f" Token: {original_tokens} -> {final_tokens} (压缩率: {final_tokens/original_tokens:.1%})")
|
||||
# L6:压缩块全空时 original_tokens 可能为 0,除零防护
|
||||
print(f" Token: {original_tokens} -> {final_tokens} (压缩率: {final_tokens / max(original_tokens, 1):.1%})")
|
||||
|
||||
# 检查是否压缩到阈值以下
|
||||
if final_count >= threshold:
|
||||
|
||||
@@ -15,12 +15,14 @@ subsystems.conversation_history - 对话历史管理器
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional
|
||||
|
||||
from common import ConfigManager, Logger, constants
|
||||
from common.atomic import atomic_write_text
|
||||
|
||||
if TYPE_CHECKING:
|
||||
# 仅用于类型注解,避免运行时循环导入;MemoryDatabase 来自同包 subsystems.db
|
||||
@@ -68,10 +70,17 @@ class ConversationHistory:
|
||||
return {"conversations": [], "last_id": 0}
|
||||
return {"conversations": [], "last_id": 0}
|
||||
|
||||
@staticmethod
|
||||
def _atomic_write_json(path: Path, data) -> None:
|
||||
"""原子写 JSON(I 批,审计 2026-08-18):先写同目录临时文件再
|
||||
os.replace——此前 'w' 直写,进程中断/崩溃留半截 JSON(current 损坏
|
||||
后仅静默置 None,历史内容丢失)。J5(2026-08-19):委托 common.atomic
|
||||
(mkstemp 随机名),修复旧实现固定 .tmp 名在多进程并发写同路径互踩。"""
|
||||
atomic_write_text(str(path), json.dumps(data, ensure_ascii=False, indent=2))
|
||||
|
||||
def _save_history_index(self):
|
||||
"""保存历史索引"""
|
||||
with open(self.history_file, 'w', encoding='utf-8') as f:
|
||||
json.dump(self.history_index, f, ensure_ascii=False, indent=2)
|
||||
self._atomic_write_json(self.history_file, self.history_index)
|
||||
|
||||
def _load_current_conversation(self):
|
||||
"""加载当前的对话(如果有)"""
|
||||
@@ -86,8 +95,7 @@ class ConversationHistory:
|
||||
def save_current_conversation(self):
|
||||
"""保存当前对话"""
|
||||
if self.current_conversation:
|
||||
with open(self.current_file, 'w', encoding='utf-8') as f:
|
||||
json.dump(self.current_conversation, f, ensure_ascii=False, indent=2)
|
||||
self._atomic_write_json(self.current_file, self.current_conversation)
|
||||
|
||||
def start_new_conversation(self, topic: str = None) -> Dict:
|
||||
"""开始新对话"""
|
||||
@@ -145,8 +153,7 @@ class ConversationHistory:
|
||||
counter += 1
|
||||
archive_path = self.archive_dir / archive_filename
|
||||
|
||||
with open(archive_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(self.current_conversation, f, ensure_ascii=False, indent=2)
|
||||
self._atomic_write_json(archive_path, self.current_conversation)
|
||||
|
||||
self.history_index["conversations"].append({
|
||||
"id": self.current_conversation["id"],
|
||||
@@ -350,7 +357,13 @@ class ConversationHistory:
|
||||
if file_name == constants.CURRENT_CONVERSATION_FILE and self.current_conversation:
|
||||
return self.current_conversation
|
||||
|
||||
file_path = self.archive_dir / file_name
|
||||
# I 批(审计 2026-08-18):file_name 目前无 LLM 可控调用面,但一旦
|
||||
# 接入即成为任意 .json 读取通道——realpath 解析软链后再做目录前缀校验
|
||||
archive_root = os.path.normcase(os.path.realpath(str(self.archive_dir)))
|
||||
target = os.path.normcase(os.path.realpath(str(self.archive_dir / file_name)))
|
||||
if not target.startswith(archive_root + os.sep):
|
||||
return None
|
||||
file_path = Path(target)
|
||||
if file_path.exists():
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
return json.load(f)
|
||||
|
||||
+149
-47
@@ -33,9 +33,18 @@ class MemoryDatabase:
|
||||
self.logger = Logger("MemoryDB")
|
||||
self._init_database()
|
||||
|
||||
def _connect(self) -> sqlite3.Connection:
|
||||
"""统一连接工厂(T25,审计 2026-08-21):对齐 MemoDatabase J9 修复——
|
||||
busy_timeout 覆盖全部路径。两类默认共用同一 memories.db,此前
|
||||
MemoryDatabase 直连只靠 sqlite3 默认 5s 等待,CLI 与服务进程并发
|
||||
写时锁冲突即抛 database is locked(记忆保存/检索失败)。"""
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
conn.execute("PRAGMA busy_timeout = 10000")
|
||||
return conn
|
||||
|
||||
def _init_database(self):
|
||||
"""初始化数据库表"""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
with self._connect() as conn:
|
||||
# WAL 模式:读写并发不互斥,缓解多 CLI 进程并发写锁等待
|
||||
conn.execute("PRAGMA journal_mode=WAL")
|
||||
cursor = conn.cursor()
|
||||
@@ -118,7 +127,7 @@ class MemoryDatabase:
|
||||
tags_str = ",".join(tags) if tags else ""
|
||||
created_at = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
with self._connect() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
INSERT INTO memories (topic, summary, content, start_time, end_time, created_at, tags, importance)
|
||||
@@ -139,13 +148,16 @@ class MemoryDatabase:
|
||||
|
||||
def delete_memory(self, memory_id: int) -> bool:
|
||||
"""删除记忆"""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
with self._connect() as conn:
|
||||
cursor = conn.cursor()
|
||||
# J3(审计 2026-08-19):必须先删 fts 再删内容表——外部内容表
|
||||
# (content=memories)删除索引条目需从内容表回读旧值,内容行
|
||||
# 先删会导致索引条目永久残留(MATCH 仍可命中已删 rowid)
|
||||
cursor.execute('DELETE FROM memories_fts WHERE rowid = ?', (memory_id,))
|
||||
cursor.execute('DELETE FROM memories WHERE id = ?', (memory_id,))
|
||||
deleted = cursor.rowcount > 0
|
||||
|
||||
if deleted:
|
||||
cursor.execute('DELETE FROM memories_fts WHERE rowid = ?', (memory_id,))
|
||||
conn.commit()
|
||||
self.logger.info(f"已删除记忆 ID: {memory_id}")
|
||||
|
||||
@@ -163,7 +175,7 @@ class MemoryDatabase:
|
||||
if limit is None:
|
||||
limit = constants.MEMORY_MAX_SEARCH_RESULTS
|
||||
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
with self._connect() as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
|
||||
@@ -232,14 +244,15 @@ class MemoryDatabase:
|
||||
return results
|
||||
|
||||
def get_recent_memories(self, limit: int = 10) -> List[Dict]:
|
||||
"""获取最近的记忆。limit=None 时不限制返回条数(供分页使用)。"""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
"""获取最近的记忆。limit=None 时使用默认上限(M4:此前 None 无 LIMIT
|
||||
全表加载,记忆库增长后单次调用 fetchall 全量入内存;与 search_memories
|
||||
的 None→默认上限语义对齐)。"""
|
||||
if limit is None:
|
||||
limit = constants.MEMORY_MAX_SEARCH_RESULTS
|
||||
with self._connect() as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
if limit is None:
|
||||
cursor.execute('SELECT * FROM memories ORDER BY end_time DESC')
|
||||
else:
|
||||
cursor.execute('SELECT * FROM memories ORDER BY end_time DESC LIMIT ?', (limit,))
|
||||
cursor.execute('SELECT * FROM memories ORDER BY end_time DESC LIMIT ?', (limit,))
|
||||
|
||||
results = []
|
||||
for row in cursor.fetchall():
|
||||
@@ -252,7 +265,7 @@ class MemoryDatabase:
|
||||
|
||||
def get_memory_by_id(self, memory_id: int) -> Optional[Dict]:
|
||||
"""根据ID获取记忆"""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
with self._connect() as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('SELECT * FROM memories WHERE id = ?', (memory_id,))
|
||||
@@ -270,14 +283,15 @@ class MemoryDatabase:
|
||||
if not memory_ids:
|
||||
return 0
|
||||
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
with self._connect() as conn:
|
||||
cursor = conn.cursor()
|
||||
placeholders = ','.join(['?'] * len(memory_ids))
|
||||
# J3:同 delete_memory——先删 fts 再删内容表,防索引条目残留
|
||||
cursor.execute(f'DELETE FROM memories_fts WHERE rowid IN ({placeholders})', memory_ids)
|
||||
cursor.execute(f'DELETE FROM memories WHERE id IN ({placeholders})', memory_ids)
|
||||
deleted_count = cursor.rowcount
|
||||
|
||||
if deleted_count > 0:
|
||||
cursor.execute(f'DELETE FROM memories_fts WHERE rowid IN ({placeholders})', memory_ids)
|
||||
conn.commit()
|
||||
self.logger.info(f"已批量删除 {deleted_count} 条记忆")
|
||||
|
||||
@@ -299,10 +313,19 @@ class MemoDatabase:
|
||||
self.logger = Logger("MemoDB")
|
||||
self._init_database()
|
||||
|
||||
def _connect(self) -> sqlite3.Connection:
|
||||
"""统一连接工厂(J9,审计 2026-08-19):busy_timeout 覆盖全部写路径,
|
||||
不再只有 mark_as_triggered 一条路径等待锁释放"""
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
conn.execute("PRAGMA busy_timeout = 10000")
|
||||
return conn
|
||||
|
||||
def _init_database(self):
|
||||
"""初始化备忘录表"""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
with self._connect() as conn:
|
||||
cursor = conn.cursor()
|
||||
# J9:WAL 显式设置——与 MemoryDatabase 同库文件时不再依赖创建顺序
|
||||
cursor.execute("PRAGMA journal_mode=WAL")
|
||||
|
||||
cursor.execute('''
|
||||
CREATE TABLE IF NOT EXISTS memos (
|
||||
@@ -348,10 +371,37 @@ class MemoDatabase:
|
||||
repeat_interval_unit: str = 'days', repeat_end_time: str = None,
|
||||
is_immediate: bool = False) -> int:
|
||||
"""添加任务 - 所有任务都是对话任务,不再需要metadata和任务类型"""
|
||||
# J1(审计 2026-08-19):入库校验——畸形时间/数字字段此前原样入库,
|
||||
# 下游无防护解析点引发调度器失败热循环与新连接握手崩溃(审计报告
|
||||
# §13.1 J1)。CLI/WS/memo 工具三入口均经本方法,校验一处覆盖。
|
||||
# W12(审计 2026-08-20):空标题任务此前可入库(CLI "|x" / 调用方疏忽)
|
||||
if not isinstance(title, str) or not title.strip():
|
||||
raise ValueError("任务标题不能为空")
|
||||
if reminder_time:
|
||||
try:
|
||||
datetime.strptime(reminder_time, constants.DATETIME_FORMAT)
|
||||
except (ValueError, TypeError):
|
||||
raise ValueError(
|
||||
f"提醒时间格式非法: {reminder_time!r}(应为 {constants.DATETIME_FORMAT})")
|
||||
if repeat_end_time:
|
||||
try:
|
||||
datetime.strptime(repeat_end_time, constants.DATETIME_FORMAT)
|
||||
except (ValueError, TypeError):
|
||||
raise ValueError(
|
||||
f"重复结束时间格式非法: {repeat_end_time!r}(应为 {constants.DATETIME_FORMAT})")
|
||||
if repeat_type not in ("none", "daily", "weekly", "monthly", "custom"):
|
||||
raise ValueError(f"无效的重复类型: {repeat_type!r}")
|
||||
try:
|
||||
repeat_interval = int(repeat_interval)
|
||||
except (ValueError, TypeError):
|
||||
raise ValueError(f"重复间隔必须是整数: {repeat_interval!r}")
|
||||
if repeat_type == "custom" and repeat_interval <= 0:
|
||||
raise ValueError("自定义重复需要设置有效的 repeat_interval (>0)")
|
||||
|
||||
created_at = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
tags_str = ",".join(tags) if tags else ""
|
||||
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
with self._connect() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
INSERT INTO memos
|
||||
@@ -376,7 +426,7 @@ class MemoDatabase:
|
||||
|
||||
def get_memo(self, memo_id: int) -> Optional[Dict]:
|
||||
"""获取单个任务"""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
with self._connect() as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('SELECT * FROM memos WHERE id = ?', (memo_id,))
|
||||
@@ -391,7 +441,7 @@ class MemoDatabase:
|
||||
|
||||
def get_pending_memos(self, limit: int = None, before_time: str = None, include_immediate: bool = False) -> List[Dict]:
|
||||
"""获取待执行的任务 - 严格使用当前时间比较"""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
with self._connect() as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
|
||||
@@ -401,11 +451,14 @@ class MemoDatabase:
|
||||
# 如果传入了 before_time,使用它,否则使用当前时间
|
||||
compare_time = before_time if before_time else now
|
||||
|
||||
# I 批(审计 2026-08-18):去掉 SQL 层 reminder_time <= ? 预过滤——
|
||||
# 字符串比较会把格式异常的时间(如 "不是时间")静默排除在 SQL 层,
|
||||
# Python 二次验证永远看不到它们(Q5:格式异常任务按到期返回)。
|
||||
# 时间过滤改由下方 Python 解析做权威判断(本地小库全扫描可忽略)。
|
||||
query = '''
|
||||
SELECT * FROM memos
|
||||
WHERE is_completed = 0
|
||||
AND reminder_time IS NOT NULL
|
||||
AND reminder_time <= ? -- 严格小于等于比较时间
|
||||
AND (repeat_end_time IS NULL OR repeat_end_time >= ?)
|
||||
'''
|
||||
|
||||
@@ -414,11 +467,13 @@ class MemoDatabase:
|
||||
|
||||
query += ' ORDER BY reminder_time ASC, priority DESC'
|
||||
|
||||
params = [compare_time, compare_time]
|
||||
params = [compare_time]
|
||||
|
||||
if limit:
|
||||
# I 批:limit=0 语义修正——此前 `if limit:` 使 0 等同"不限",
|
||||
# 与 None=不限的惯例冲突;改为 None 才不加 LIMIT
|
||||
if limit is not None:
|
||||
query += ' LIMIT ?'
|
||||
params.append(int(limit))
|
||||
params.append(max(0, int(limit)))
|
||||
|
||||
cursor.execute(query, params)
|
||||
|
||||
@@ -439,7 +494,11 @@ class MemoDatabase:
|
||||
# 调试日志
|
||||
self.logger.debug(f"过滤掉未到期任务 #{memo['id']}: {memo['reminder_time']} > {compare_time}")
|
||||
except (ValueError, TypeError) as e:
|
||||
self.logger.warning(f"解析任务 #{memo['id']} 的提醒时间失败: {e}")
|
||||
# I 批(审计 2026-08-18):格式异常此前被静默丢弃——
|
||||
# 任务永不执行且用户无感知。改为按"到期"返回并告警,
|
||||
# 至少保证任务被处理而非永久丢失
|
||||
self.logger.warning(f"解析任务 #{memo['id']} 的提醒时间失败,按到期处理: {e}")
|
||||
results.append(memo)
|
||||
else:
|
||||
# 没有提醒时间的任务不应该在待执行列表中
|
||||
self.logger.warning(f"任务 #{memo['id']} 没有提醒时间,但被返回为待执行")
|
||||
@@ -448,7 +507,7 @@ class MemoDatabase:
|
||||
|
||||
def get_immediate_memos(self, limit: int = None) -> List[Dict]:
|
||||
"""获取待执行的即时任务"""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
with self._connect() as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
|
||||
@@ -460,9 +519,11 @@ class MemoDatabase:
|
||||
'''
|
||||
|
||||
params = []
|
||||
if limit:
|
||||
# I 批:limit=0 语义修正——此前 `if limit:` 使 0 等同"不限",
|
||||
# 与 None=不限的惯例冲突;改为 None 才不加 LIMIT
|
||||
if limit is not None:
|
||||
query += ' LIMIT ?'
|
||||
params.append(int(limit))
|
||||
params.append(max(0, int(limit)))
|
||||
|
||||
cursor.execute(query, params)
|
||||
|
||||
@@ -482,7 +543,7 @@ class MemoDatabase:
|
||||
|
||||
def get_all_memos(self, include_completed: bool = False, limit: int = 50) -> List[Dict]:
|
||||
"""获取所有任务"""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
with self._connect() as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
|
||||
@@ -521,9 +582,11 @@ class MemoDatabase:
|
||||
now = datetime.now()
|
||||
next_reminder = None
|
||||
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
with self._connect() as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
# J9:busy_timeout 已由 _connect() 工厂统一设置(全部写路径一致),
|
||||
# 多进程同库(CLI --memo-add + 服务调度)写竞争等待锁释放而非立即失败
|
||||
cursor.execute("BEGIN IMMEDIATE")
|
||||
try:
|
||||
cursor.execute('SELECT * FROM memos WHERE id = ?', (memo_id,))
|
||||
@@ -592,7 +655,11 @@ class MemoDatabase:
|
||||
def _should_end_repeat(self, memo: Dict, current_time: datetime) -> bool:
|
||||
"""检查是否应该结束重复"""
|
||||
if memo['repeat_end_time']:
|
||||
end_time = datetime.strptime(memo['repeat_end_time'], constants.DATETIME_FORMAT)
|
||||
try:
|
||||
end_time = datetime.strptime(memo['repeat_end_time'], constants.DATETIME_FORMAT)
|
||||
except (ValueError, TypeError):
|
||||
# J1:存量脏数据防御——解析失败视为无结束时间,不打断触发事务
|
||||
return False
|
||||
if current_time > end_time:
|
||||
return True
|
||||
return False
|
||||
@@ -622,24 +689,37 @@ class MemoDatabase:
|
||||
next_time = current_time.replace(year=current_time.year + 1, month=1, day=1) - timedelta(days=1)
|
||||
else:
|
||||
next_time = current_time.replace(month=current_time.month + 1, day=1) - timedelta(days=1)
|
||||
elif memo['repeat_type'] == 'custom' and memo['repeat_interval'] > 0:
|
||||
unit = memo.get('repeat_interval_unit', 'days')
|
||||
if unit == 'minutes':
|
||||
# 每N分钟重复:当前时间 + N分钟
|
||||
next_time = current_time + timedelta(minutes=memo['repeat_interval'])
|
||||
elif unit == 'hours':
|
||||
# 每N小时重复:当前时间 + N小时
|
||||
next_time = current_time + timedelta(hours=memo['repeat_interval'])
|
||||
else: # days
|
||||
# 每N天重复:当前时间 + N天
|
||||
next_time = current_time + timedelta(days=memo['repeat_interval'])
|
||||
elif memo['repeat_type'] == 'custom':
|
||||
# J1:repeat_interval 类型异常(存量脏数据)按 0 处理,
|
||||
# 不再抛 TypeError 打崩触发事务引发热循环
|
||||
try:
|
||||
repeat_interval = int(memo.get('repeat_interval') or 0)
|
||||
except (ValueError, TypeError):
|
||||
repeat_interval = 0
|
||||
if repeat_interval > 0:
|
||||
unit = memo.get('repeat_interval_unit', 'days')
|
||||
if unit == 'minutes':
|
||||
# 每N分钟重复:当前时间 + N分钟
|
||||
next_time = current_time + timedelta(minutes=repeat_interval)
|
||||
elif unit == 'hours':
|
||||
# 每N小时重复:当前时间 + N小时
|
||||
next_time = current_time + timedelta(hours=repeat_interval)
|
||||
else: # days
|
||||
# 每N天重复:当前时间 + N天
|
||||
next_time = current_time + timedelta(days=repeat_interval)
|
||||
else:
|
||||
return None
|
||||
else:
|
||||
return None
|
||||
|
||||
# 检查是否超过重复结束时间
|
||||
if memo['repeat_end_time']:
|
||||
end_time = datetime.strptime(memo['repeat_end_time'], constants.DATETIME_FORMAT)
|
||||
if next_time > end_time:
|
||||
try:
|
||||
end_time = datetime.strptime(memo['repeat_end_time'], constants.DATETIME_FORMAT)
|
||||
except (ValueError, TypeError):
|
||||
# J1:存量脏数据防御——解析失败视为无结束时间
|
||||
end_time = None
|
||||
if end_time and next_time > end_time:
|
||||
return None
|
||||
|
||||
return next_time
|
||||
@@ -656,7 +736,7 @@ class MemoDatabase:
|
||||
|
||||
completed_at = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
with self._connect() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
UPDATE memos
|
||||
@@ -672,9 +752,18 @@ class MemoDatabase:
|
||||
return True
|
||||
return False
|
||||
|
||||
def delete_memo(self, memo_id: int) -> bool:
|
||||
"""物理删除任务(M5:delete 语义此前被 complete_memo 顶替——只标记
|
||||
完成不删行,重复任务还会被推进而非移除,stats 计数失真)"""
|
||||
with self._connect() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("DELETE FROM memos WHERE id = ?", (memo_id,))
|
||||
conn.commit()
|
||||
return cursor.rowcount > 0
|
||||
|
||||
def get_memo_stats(self) -> Dict:
|
||||
"""获取任务统计信息"""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
with self._connect() as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute('SELECT COUNT(*) FROM memos')
|
||||
@@ -707,8 +796,9 @@ class MemoDatabase:
|
||||
}
|
||||
|
||||
def add_user_task(self, title: str, content: str = None, reminder_time: str = None,
|
||||
repeat_type: str = 'none', repeat_interval: int = 0,
|
||||
repeat_interval_unit: str = 'days', is_immediate: bool = False) -> int:
|
||||
repeat_type: str = 'none', repeat_interval: int = 0,
|
||||
repeat_interval_unit: str = 'days', repeat_end_time: str = None,
|
||||
is_immediate: bool = False) -> int:
|
||||
"""添加用户任务(前端直接调用的任务)- 所有任务都是对话任务"""
|
||||
if not constants.MEMO_ALLOW_USER_TASKS:
|
||||
raise Exception("用户任务功能已禁用")
|
||||
@@ -720,11 +810,23 @@ class MemoDatabase:
|
||||
repeat_type=repeat_type,
|
||||
repeat_interval=repeat_interval,
|
||||
repeat_interval_unit=repeat_interval_unit,
|
||||
repeat_end_time=repeat_end_time,
|
||||
created_by="user",
|
||||
source="websocket",
|
||||
is_immediate=is_immediate
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_overdue_time(reminder_time: str, now: datetime) -> bool:
|
||||
"""J1:展示路径时间解析防御——格式异常按到期处理(对齐 Q5 语义),
|
||||
不再抛 ValueError 打崩新连接握手流程"""
|
||||
if not reminder_time:
|
||||
return False
|
||||
try:
|
||||
return datetime.strptime(reminder_time, constants.DATETIME_FORMAT) < now
|
||||
except (ValueError, TypeError):
|
||||
return True
|
||||
|
||||
def get_pending_tasks_for_user(self) -> List[Dict]:
|
||||
"""获取待执行任务列表(用于前端显示)"""
|
||||
pending = self.get_pending_memos(include_immediate=False)
|
||||
@@ -739,7 +841,7 @@ class MemoDatabase:
|
||||
"content": memo['content'],
|
||||
"reminder_time": memo['reminder_time'],
|
||||
"priority": memo['priority'],
|
||||
"is_overdue": datetime.strptime(memo['reminder_time'], constants.DATETIME_FORMAT) < now if memo['reminder_time'] else False,
|
||||
"is_overdue": self._is_overdue_time(memo['reminder_time'], now),
|
||||
"repeat_type": memo['repeat_type'],
|
||||
"repeat_interval": memo.get('repeat_interval', 0),
|
||||
"repeat_interval_unit": memo.get('repeat_interval_unit', 'days'),
|
||||
|
||||
+189
-41
@@ -13,6 +13,7 @@ subsystems.executors - 命令执行器
|
||||
|
||||
import os
|
||||
import queue
|
||||
import signal
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
@@ -21,6 +22,36 @@ from typing import Dict, Optional, Tuple
|
||||
|
||||
from common import ConfigManager, constants, Logger
|
||||
|
||||
# T21(审计 2026-08-21):env 覆盖黑名单——工具 schema 中 env 是开放
|
||||
# object,LLM(可被网页内容提示注入诱导)可传 LD_PRELOAD/PYTHONPATH/
|
||||
# PATH 等在 shell=True 之上进一步劫持进程加载(POSIX LD_PRELOAD 可加载
|
||||
# 任意共享库)。命中黑名单的覆盖键拒绝并告警
|
||||
ENV_DENY_EXACT = frozenset({
|
||||
"PATH", "PYTHONPATH", "PYTHONHOME", "SHELL", "IFS",
|
||||
"LD_LIBRARY_PATH", "LD_PRELOAD", "LD_AUDIT",
|
||||
"DYLD_LIBRARY_PATH", "DYLD_INSERT_LIBRARIES",
|
||||
"PSMODULEPATH", "COMSPEC",
|
||||
})
|
||||
ENV_DENY_PREFIXES = ("LD_", "DYLD_")
|
||||
|
||||
|
||||
def filter_env_overrides(env) -> Tuple[Dict, list]:
|
||||
"""T21:过滤 env 覆盖项,返回 (安全覆盖, 被拒绝键列表)。
|
||||
|
||||
键大小写不敏感匹配(Windows 环境变量不区分大小写);
|
||||
LD_*/DYLD_* 前缀整体拒绝(加载器劫持面)。
|
||||
"""
|
||||
if not env:
|
||||
return {}, []
|
||||
safe, rejected = {}, []
|
||||
for key, value in dict(env).items():
|
||||
ku = str(key).upper()
|
||||
if ku in ENV_DENY_EXACT or ku.startswith(ENV_DENY_PREFIXES):
|
||||
rejected.append(key)
|
||||
else:
|
||||
safe[key] = value
|
||||
return safe, rejected
|
||||
|
||||
|
||||
class CommandExecutor(ABC):
|
||||
"""命令执行器抽象基类 - 负责执行系统命令"""
|
||||
@@ -70,22 +101,104 @@ class SubprocessCommandExecutor(CommandExecutor):
|
||||
|
||||
def __init__(self, logger=None):
|
||||
super().__init__(logger)
|
||||
self._current_process = None
|
||||
self._current_process = None # 兼容保留:最近一次注册的进程(诊断用)
|
||||
self._current_pgid = None # L9:Popen 后立即捕获,shell 死亡后 getpgid 会失败
|
||||
self._process_lock = threading.Lock()
|
||||
# T5:单槽 _current_process 在并发执行时互相覆盖——先启动的执行
|
||||
# 超时/取消会误杀后来者的进程组(killpg 连坐孙进程),先结束的
|
||||
# finally 清引用导致后来者永远无法取消。改为令牌多槽注册表,
|
||||
# 终止只作用于自己槽位内的 (proc, pgid)。
|
||||
self._proc_slots: Dict[int, Tuple] = {}
|
||||
self._token_seq = 0
|
||||
|
||||
def cancel(self):
|
||||
"""取消当前正在运行的命令(kill子进程)。线程安全。"""
|
||||
def _register_proc(self, process):
|
||||
"""T5:注册本次执行的进程槽,返回 (token, pgid)"""
|
||||
with self._process_lock:
|
||||
proc = self._current_process
|
||||
if proc is not None:
|
||||
self._token_seq += 1
|
||||
token = self._token_seq
|
||||
pgid = None
|
||||
if os.name == "posix":
|
||||
try:
|
||||
pgid = os.getpgid(process.pid)
|
||||
except Exception:
|
||||
pgid = None
|
||||
self._proc_slots[token] = (process, pgid)
|
||||
# 兼容属性指向最新注册(诊断/旧测试观察用,不参与终止逻辑)
|
||||
self._current_process = process
|
||||
self._current_pgid = pgid
|
||||
return token, pgid
|
||||
|
||||
def _unregister_proc(self, token):
|
||||
"""T5:仅移除自己的槽位;全部槽位清空时同步清兼容属性"""
|
||||
with self._process_lock:
|
||||
self._proc_slots.pop(token, None)
|
||||
if not self._proc_slots:
|
||||
self._current_process = None
|
||||
self._current_pgid = None
|
||||
|
||||
def _terminate_tree(self, proc, force: bool = False, pgid=None):
|
||||
"""终止整个进程组(L9:POSIX 下 shell=True 只杀 shell 时,孙进程
|
||||
(`(x &)` 等)仍持管道写端——reader 阻塞无 EOF、fd/线程悬挂至孙
|
||||
进程退出,写满 64KB 管道缓冲后孙进程也卡死。start_new_session 使
|
||||
shell 成为会话首进程,按组信号可一次覆盖整棵进程树)。
|
||||
组 ID 在 Popen 后立即保存:shell 死亡后 os.getpgid(pid) 会抛
|
||||
ProcessLookupError,届时无法再从进程反查组。
|
||||
T5:pgid 由调用方显式传入(注册槽位时捕获),不再读共享
|
||||
_current_pgid——并发执行时那是别人的进程组。
|
||||
T11(审计 2026-08-21):Windows 无进程组概念,单进程 terminate/
|
||||
kill 收割不到 shell=True 生成的 cmd.exe /c 之下 spawn 的孙进程
|
||||
(超时/取消后孙进程成孤儿,持续占用 CPU/端口/文件锁)。用
|
||||
taskkill /T 递归终止整棵进程树;Windows 无 SIGTERM 等价物,
|
||||
taskkill 不带 /F 仅对有消息窗口的 GUI 进程有效,控制台进程
|
||||
杀不掉,故统一 /F 强制,force 参数在 Windows 上不区分。"""
|
||||
try:
|
||||
if os.name == "posix":
|
||||
sig = signal.SIGKILL if force else signal.SIGTERM
|
||||
if pgid is None:
|
||||
try:
|
||||
pgid = os.getpgid(proc.pid)
|
||||
except Exception:
|
||||
pgid = None
|
||||
if pgid is not None:
|
||||
os.killpg(pgid, sig)
|
||||
else:
|
||||
proc.kill() if force else proc.terminate()
|
||||
else:
|
||||
# T11:Windows 树收割——taskkill /T /F 递归终止子孙进程,
|
||||
# 失败(进程已退出/无权限)时兜底单进程 kill
|
||||
try:
|
||||
subprocess.run(
|
||||
["taskkill", "/PID", str(proc.pid), "/T", "/F"],
|
||||
capture_output=True, timeout=10)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
if proc.poll() is None:
|
||||
proc.kill()
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
try:
|
||||
if proc.poll() is None:
|
||||
proc.kill()
|
||||
self.logger.info("已终止正在运行的子进程")
|
||||
proc.kill()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _read_with_limit(self, process, timeout: int, start_time: float, output_queue=None) -> Tuple[str, bool, bool]:
|
||||
def cancel(self):
|
||||
"""取消本执行器当前注册的全部运行中命令(kill 子进程组)。线程安全。
|
||||
T5:多槽令牌注册表——并发执行时定向终止各自槽位,互不误杀;
|
||||
无参调用语义为"终止本执行器全部在途命令"(中断/关闭路径)。"""
|
||||
with self._process_lock:
|
||||
slots = list(self._proc_slots.values())
|
||||
for proc, pgid in slots:
|
||||
try:
|
||||
if proc.poll() is None:
|
||||
self._terminate_tree(proc, force=True, pgid=pgid)
|
||||
self.logger.info("已终止正在运行的子进程组")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _read_with_limit(self, process, timeout: int, start_time: float, output_queue=None,
|
||||
pgid=None) -> Tuple[str, bool, bool]:
|
||||
"""从进程读取输出,限制总大小和超时。返回 (output_str, timeout_occurred, truncated)。
|
||||
|
||||
H2 修复:阻塞 readline 在独立 reader 线程执行,主循环带超时轮询队列。
|
||||
@@ -116,15 +229,17 @@ class SubprocessCommandExecutor(CommandExecutor):
|
||||
shell_dead_since = None # shell 进程已死但 reader 未收到 EOF 的起始时刻
|
||||
while True:
|
||||
# 检查超时(readline 阻塞不再影响本判断)
|
||||
# J16(审计 2026-08-19):timeout<=0 语义为"不设超时"(有意为之,
|
||||
# 调用方显式传 0 禁用);registry 层 3600s 流式兜底仍生效
|
||||
if timeout > 0 and time.time() - start_time > timeout:
|
||||
process.terminate()
|
||||
self._terminate_tree(process, pgid=pgid)
|
||||
timeout_occurred = True
|
||||
break
|
||||
|
||||
# 检查输出大小
|
||||
if total_bytes >= self.MAX_OUTPUT_BYTES:
|
||||
truncated = True
|
||||
process.terminate()
|
||||
self._terminate_tree(process, pgid=pgid)
|
||||
break
|
||||
|
||||
try:
|
||||
@@ -140,6 +255,12 @@ class SubprocessCommandExecutor(CommandExecutor):
|
||||
if shell_dead_since is None:
|
||||
shell_dead_since = time.time()
|
||||
elif time.time() - shell_dead_since > 1.0:
|
||||
# L9:shell 已死但孙进程仍持管道写端(`(sleep 300 &) ;
|
||||
# echo done` 形态)。只靠 daemon reader 兜底会悬挂线程
|
||||
# 与管道 fd 至孙进程退出,写满 64KB 缓冲后孙进程自身
|
||||
# 也卡死。此处强杀进程组收尾(nohup 重定向输出的
|
||||
# 真正脱离场景不持管道,不会走到这里)。
|
||||
self._terminate_tree(process, force=True, pgid=pgid)
|
||||
break
|
||||
continue
|
||||
|
||||
@@ -147,10 +268,12 @@ class SubprocessCommandExecutor(CommandExecutor):
|
||||
if total_bytes + line_bytes > self.MAX_OUTPUT_BYTES and not truncated:
|
||||
allowed = self.MAX_OUTPUT_BYTES - total_bytes
|
||||
if allowed > 0:
|
||||
output_chunks.append(line.encode('utf-8', errors='replace')[:allowed].decode('utf-8', errors='replace'))
|
||||
# L5:decode 用 errors='ignore'——按字节截断可能落在多字节
|
||||
# 字符中间,errors='replace' 会在行尾留下半个字符(�)
|
||||
output_chunks.append(line.encode('utf-8', errors='replace')[:allowed].decode('utf-8', errors='ignore'))
|
||||
total_bytes = self.MAX_OUTPUT_BYTES
|
||||
truncated = True
|
||||
process.terminate()
|
||||
self._terminate_tree(process, pgid=pgid)
|
||||
break
|
||||
|
||||
output_chunks.append(line)
|
||||
@@ -186,18 +309,31 @@ class SubprocessCommandExecutor(CommandExecutor):
|
||||
if not self.is_enabled():
|
||||
return "❌ 命令执行功能已禁用"
|
||||
|
||||
# T24:非流式路径没有 registry 3600s 流式兜底,timeout<=0("不设
|
||||
# 超时")会让挂起命令永久阻塞调用线程——下限夹取到 1s 快速失败。
|
||||
# 流式路径保留 0=不限语义(registry 兜底在位)
|
||||
try:
|
||||
timeout = max(1, int(timeout))
|
||||
except (TypeError, ValueError):
|
||||
timeout = 30
|
||||
|
||||
# 默认工作目录为workspace,避免在启动目录散落文件
|
||||
if work_dir is None or work_dir == ".":
|
||||
work_dir = constants.WORKSPACE_DIR
|
||||
|
||||
self.logger.info(f"执行命令: {command} (work_dir={work_dir})")
|
||||
|
||||
# 准备环境变量
|
||||
# 准备环境变量(T21:黑名单键拒绝,防加载器/解释器劫持)
|
||||
exec_env = os.environ.copy()
|
||||
if env:
|
||||
exec_env.update(env)
|
||||
safe_env, rejected_keys = filter_env_overrides(env)
|
||||
if rejected_keys:
|
||||
self.logger.warning(f"env 覆盖被黑名单拒绝(T21): {rejected_keys}")
|
||||
exec_env.update(safe_env)
|
||||
|
||||
process = None
|
||||
token = None
|
||||
pgid = None
|
||||
try:
|
||||
process = subprocess.Popen(
|
||||
command,
|
||||
@@ -208,26 +344,28 @@ class SubprocessCommandExecutor(CommandExecutor):
|
||||
text=True,
|
||||
bufsize=1,
|
||||
universal_newlines=True,
|
||||
encoding="utf-8", # T16:显式 UTF-8,避免 Windows cp936 解码 UTF-8 输出
|
||||
errors="replace", # T16:非法字节不抛异常中断 reader 线程
|
||||
env=exec_env,
|
||||
cwd=os.path.expanduser(work_dir)
|
||||
cwd=os.path.expanduser(work_dir),
|
||||
# L9:POSIX 下 shell 成为会话首进程,超时/取消可按进程组
|
||||
# 一次终止整棵进程树(含 & 后台孙进程),消除管道悬挂
|
||||
start_new_session=(os.name == "posix")
|
||||
)
|
||||
# 保存进程引用以便cancel()可以kill它
|
||||
with self._process_lock:
|
||||
self._current_process = process
|
||||
# T5:注册本执行专属槽位,终止/取消只作用于自己的 (proc, pgid)
|
||||
token, pgid = self._register_proc(process)
|
||||
|
||||
start_time = time.time()
|
||||
output_str, timeout_occurred, truncated = self._read_with_limit(process, timeout, start_time)
|
||||
output_str, timeout_occurred, truncated = self._read_with_limit(
|
||||
process, timeout, start_time, pgid=pgid)
|
||||
|
||||
# 确保进程终止
|
||||
if process.poll() is None:
|
||||
if timeout_occurred or truncated:
|
||||
process.kill()
|
||||
else:
|
||||
process.terminate()
|
||||
self._terminate_tree(process, force=timeout_occurred or truncated, pgid=pgid)
|
||||
try:
|
||||
process.wait(timeout=2)
|
||||
except:
|
||||
process.kill()
|
||||
self._terminate_tree(process, force=True, pgid=pgid)
|
||||
|
||||
return_code = process.returncode if process.returncode is not None else -1
|
||||
|
||||
@@ -255,10 +393,10 @@ class SubprocessCommandExecutor(CommandExecutor):
|
||||
return f"❌ 执行错误: {str(e)}"
|
||||
finally:
|
||||
if process and process.poll() is None:
|
||||
process.kill()
|
||||
# 清除当前进程引用
|
||||
with self._process_lock:
|
||||
self._current_process = None
|
||||
self._terminate_tree(process, force=True, pgid=pgid)
|
||||
# T5:仅注销自己的槽位,不再无差别清空共享引用
|
||||
if token is not None:
|
||||
self._unregister_proc(token)
|
||||
|
||||
def execute_streaming(self, command: str, output_queue: queue.Queue, timeout: int = 30,
|
||||
work_dir: str = None, env: Dict = None, **kwargs):
|
||||
@@ -279,12 +417,17 @@ class SubprocessCommandExecutor(CommandExecutor):
|
||||
output_queue.put(("line", f"📁 工作目录: {work_dir}\n"))
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
# 准备环境变量
|
||||
# 准备环境变量(T21:黑名单键拒绝,防加载器/解释器劫持)
|
||||
exec_env = os.environ.copy()
|
||||
if env:
|
||||
exec_env.update(env)
|
||||
safe_env, rejected_keys = filter_env_overrides(env)
|
||||
if rejected_keys:
|
||||
self.logger.warning(f"env 覆盖被黑名单拒绝(T21): {rejected_keys}")
|
||||
exec_env.update(safe_env)
|
||||
|
||||
process = None
|
||||
token = None
|
||||
pgid = None
|
||||
try:
|
||||
process = subprocess.Popen(
|
||||
command,
|
||||
@@ -295,18 +438,22 @@ class SubprocessCommandExecutor(CommandExecutor):
|
||||
text=True,
|
||||
bufsize=1,
|
||||
universal_newlines=True,
|
||||
encoding="utf-8", # T16:显式 UTF-8,避免 Windows cp936 解码 UTF-8 输出
|
||||
errors="replace", # T16:非法字节不抛异常中断 reader 线程
|
||||
env=exec_env,
|
||||
cwd=os.path.expanduser(work_dir)
|
||||
cwd=os.path.expanduser(work_dir),
|
||||
# L9:POSIX 下 shell 成为会话首进程,超时/取消可按进程组
|
||||
# 一次终止整棵进程树(含 & 后台孙进程),消除管道悬挂
|
||||
start_new_session=(os.name == "posix")
|
||||
)
|
||||
self.process = process
|
||||
# 保存进程引用以便cancel()可以kill它
|
||||
with self._process_lock:
|
||||
self._current_process = process
|
||||
# T5:注册本执行专属槽位,终止/取消只作用于自己的 (proc, pgid)
|
||||
token, pgid = self._register_proc(process)
|
||||
|
||||
start_time = time.time()
|
||||
# H2 修复:复用线程化读取(reader 线程 + 超时轮询),行经 output_queue 流式转发
|
||||
output_str, timeout_occurred, truncated = self._read_with_limit(
|
||||
process, timeout, start_time, output_queue=output_queue)
|
||||
process, timeout, start_time, output_queue=output_queue, pgid=pgid)
|
||||
|
||||
if timeout_occurred:
|
||||
output_queue.put(("line", f"\n⏰ 命令执行超时({timeout}秒)\n"))
|
||||
@@ -315,7 +462,7 @@ class SubprocessCommandExecutor(CommandExecutor):
|
||||
|
||||
# 如果超时或截断,确保进程被终止
|
||||
if (timeout_occurred or truncated) and process.poll() is None:
|
||||
process.kill()
|
||||
self._terminate_tree(process, force=True, pgid=pgid)
|
||||
try:
|
||||
process.wait(timeout=2)
|
||||
except:
|
||||
@@ -352,11 +499,12 @@ class SubprocessCommandExecutor(CommandExecutor):
|
||||
output_queue.put(("complete", f"❌ 执行错误: {str(e)}"))
|
||||
finally:
|
||||
if process and process.poll() is None:
|
||||
process.terminate()
|
||||
self._terminate_tree(process, pgid=pgid)
|
||||
try:
|
||||
process.wait(timeout=2)
|
||||
except:
|
||||
process.kill()
|
||||
self._terminate_tree(process, force=True, pgid=pgid)
|
||||
self.process = None
|
||||
with self._process_lock:
|
||||
self._current_process = None
|
||||
# T5:仅注销自己的槽位,不再无差别清空共享引用
|
||||
if token is not None:
|
||||
self._unregister_proc(token)
|
||||
@@ -0,0 +1,779 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
subsystems.process_manager - 后台进程管理器
|
||||
|
||||
让智能体能够启动并管理长驻后台进程(dev server / 训练任务 / watch 等),
|
||||
弥补 run_command 同步阻塞、超时即杀的不足:
|
||||
- bg 语义:start() 立即返回句柄,进程独立于对话轮次存活
|
||||
- 有界环形输出缓冲:每进程上限 buffer_per_process(默认 256KB),
|
||||
解决 nohup+日志文件方案输出无界增长的问题
|
||||
- 注册表:自增句柄跨会话/跨子代理可见(全局单例)
|
||||
- 生命周期:惰性状态收割(reader 线程 EOF 即 poll 收割防僵尸)+
|
||||
TTL 惰性清理 + 退出统一收尾(非 detach 进程默认全杀)
|
||||
|
||||
安全边界(对齐 SubprocessCommandExecutor 既有惯例):
|
||||
- 输出单行截断 16KB,防超长单行撑爆内存
|
||||
- 运行中进程数上限(默认 16)+ 总条目上限(默认 32,已退出条目
|
||||
保留尾部输出供验尸,LRU 淘汰)
|
||||
- POSIX start_new_session + 进程组信号(L9 语义:孙进程不悬挂);
|
||||
Windows 无进程组概念,回退单进程 terminate/kill
|
||||
- 中断语义:用户 Ctrl+C / 取消事件不影响后台进程(后台的本意就是
|
||||
脱离本轮对话生命周期);aibox 进程退出时经 atexit 统一收尾,
|
||||
detach=true 的进程豁免
|
||||
|
||||
依赖:common (ConfigManager, constants, Logger) + 标准库
|
||||
不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
import atexit
|
||||
import collections
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from common import ConfigManager, constants, Logger
|
||||
from subsystems.executors import filter_env_overrides
|
||||
|
||||
# 单行输出最大字节数,超出截断(防超长单行撑爆环形缓冲内存预算)
|
||||
MAX_LINE_BYTES = 16 * 1024
|
||||
|
||||
# 单次 bg_output 拉取的硬上限(对齐执行器 1MB 输出限制)
|
||||
MAX_FETCH_BYTES = 1 * 1024 * 1024
|
||||
|
||||
|
||||
class _ProcEntry:
|
||||
"""单个后台进程的注册表条目(内部类)"""
|
||||
|
||||
__slots__ = (
|
||||
"handle", "name", "command", "popen", "pgid", "started_at",
|
||||
"detach", "max_lifetime", "owner", "status", "exit_code", "end_time",
|
||||
"lines", "buffer_bytes", "max_buffer", "next_seq", "evicted",
|
||||
"lock", "reader_done", "reader", "stdin_lock", "log_fh",
|
||||
)
|
||||
|
||||
# 状态取值:running / exited(自然退出)/ killed(被管理器终止)/ expired(TTL 到期)
|
||||
def __init__(self, handle: int, name: str, command: str, popen, pgid,
|
||||
detach: bool, max_lifetime: float, buffer_bytes: int,
|
||||
owner: Optional[str] = None):
|
||||
self.handle = handle
|
||||
self.name = name
|
||||
self.command = command
|
||||
self.popen = popen
|
||||
self.pgid = pgid
|
||||
self.started_at = time.time()
|
||||
self.detach = detach
|
||||
self.max_lifetime = max_lifetime
|
||||
self.owner = owner
|
||||
self.status = "running"
|
||||
self.exit_code = None
|
||||
self.end_time = None
|
||||
# 环形输出缓冲:deque of (seq, line, nbytes),seq 单调递增不回收,
|
||||
# 淘汰最旧行仍保持 seq 连续性——bg_output 的 since 游标据此续读
|
||||
self.lines = collections.deque()
|
||||
self.buffer_bytes = 0
|
||||
self.max_buffer = max(1024, int(buffer_bytes))
|
||||
self.next_seq = 1
|
||||
self.evicted = 0
|
||||
self.lock = threading.RLock() # T8:可重入,write_input 持锁期间可调用 _record_line
|
||||
self.reader_done = threading.Event()
|
||||
self.reader = None
|
||||
# 交互写入:单写锁保序(bg_write 串行化,防交错输入)
|
||||
self.stdin_lock = threading.Lock()
|
||||
# 可选全档镜像文件句柄(log_file=True 时由 reader 线程 tee 写入)
|
||||
self.log_fh = None
|
||||
|
||||
|
||||
class ProcessManager:
|
||||
"""后台进程管理器:启动 / 列举 / 增量拉取输出 / 终止 / 退出收尾"""
|
||||
|
||||
def __init__(self, logger: Optional[Logger] = None,
|
||||
max_processes: int = None, max_entries: int = None,
|
||||
buffer_per_process: int = None, max_lifetime: float = None):
|
||||
self.logger = logger or Logger("ProcessManager")
|
||||
self.config = ConfigManager()
|
||||
self.max_processes = int(max_processes if max_processes is not None
|
||||
else constants.BG_MAX_PROCESSES)
|
||||
self.max_entries = int(max_entries if max_entries is not None
|
||||
else constants.BG_MAX_ENTRIES)
|
||||
self.buffer_per_process = int(buffer_per_process if buffer_per_process is not None
|
||||
else constants.BG_BUFFER_PER_PROCESS)
|
||||
self.max_lifetime = float(max_lifetime if max_lifetime is not None
|
||||
else constants.BG_MAX_LIFETIME)
|
||||
self.max_write_bytes = int(constants.BG_MAX_WRITE_BYTES)
|
||||
self._procs: Dict[int, _ProcEntry] = {}
|
||||
self._next_id = 1
|
||||
self._lock = threading.Lock()
|
||||
self._shutdown_done = False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 启动
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def start(self, command: str, name: str = None, work_dir: str = None,
|
||||
env: Dict = None, detach: bool = False,
|
||||
log_file: bool = False, owner: str = None) -> Dict:
|
||||
"""后台启动命令,立即返回。成功返回 {ok, handle, pid, name, log_file},
|
||||
失败返回 {ok: False, error}。
|
||||
|
||||
stdin 为管道:后续可经 write() 交互写入(行协议 REPL 问答)。
|
||||
log_file=True 时输出同步 tee 到 workspace/bg_<句柄>.log(无界
|
||||
全档,可被原生工具 tail/cat 读取;环形缓冲仍为 bg_output 视图)。"""
|
||||
if not command or not isinstance(command, str) or not command.strip():
|
||||
return {"ok": False, "error": "缺少或空的 command"}
|
||||
|
||||
# 名称规整(M5 惯例):标签仅用于展示回显,拒绝控制字符防终端/前端注入
|
||||
if name is not None:
|
||||
name = str(name).strip()
|
||||
if not name.isprintable() or len(name) > 64:
|
||||
return {"ok": False, "error": "name 含控制字符或超过 64 字符"}
|
||||
if not name:
|
||||
name = None
|
||||
if name is None:
|
||||
name = command.strip()[:40]
|
||||
|
||||
if work_dir is None or work_dir == ".":
|
||||
work_dir = constants.WORKSPACE_DIR
|
||||
|
||||
with self._lock:
|
||||
if self._shutdown_done:
|
||||
return {"ok": False, "error": "管理器已关闭"}
|
||||
# TTL 惰性清理 + 状态收割
|
||||
self._sweep_locked()
|
||||
running = [e for e in self._procs.values() if e.status == "running"]
|
||||
if len(running) >= self.max_processes:
|
||||
return {"ok": False,
|
||||
"error": f"运行中后台进程已达上限({self.max_processes}),"
|
||||
f"请先用 bg_kill 清理"}
|
||||
# 总条目上限:淘汰最老的已退出条目(保留其信息的价值最低)
|
||||
if len(self._procs) >= self.max_entries:
|
||||
exited = sorted(
|
||||
(e for e in self._procs.values() if e.status != "running"),
|
||||
key=lambda e: e.started_at)
|
||||
while self._procs and len(self._procs) >= self.max_entries and exited:
|
||||
victim = exited.pop(0)
|
||||
self._procs.pop(victim.handle, None)
|
||||
if len(self._procs) >= self.max_entries:
|
||||
return {"ok": False,
|
||||
"error": f"后台进程条目已达上限({self.max_entries})且全部运行中"}
|
||||
handle = self._next_id
|
||||
self._next_id += 1
|
||||
|
||||
exec_env = os.environ.copy()
|
||||
if env:
|
||||
# T21:黑名单键拒绝(对齐 executors,防 LD_PRELOAD/PYTHONPATH 劫持)
|
||||
safe_env, rejected_keys = filter_env_overrides(env)
|
||||
if rejected_keys:
|
||||
self.logger.warning(f"env 覆盖被黑名单拒绝(T21): {rejected_keys}")
|
||||
exec_env.update(safe_env)
|
||||
|
||||
try:
|
||||
popen = subprocess.Popen(
|
||||
command,
|
||||
shell=True,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
# 交互式 I/O:stdin 挂管道供 write() 写入。与 H3 的差异:
|
||||
# 后台进程由管理器显式管理生命周期,不存在"等待输入的
|
||||
# 命令永久挂起无人知晓"问题——不读 stdin 的进程不受影响,
|
||||
# 读 stdin 的进程获得交互能力,kill 路径三级优雅终止兜底。
|
||||
stdin=subprocess.PIPE,
|
||||
text=True,
|
||||
bufsize=1,
|
||||
universal_newlines=True,
|
||||
encoding="utf-8", # T17:显式 UTF-8,避免 Windows cp936 解码 UTF-8 输出致 reader 静默退出
|
||||
errors="replace",
|
||||
env=exec_env,
|
||||
cwd=os.path.expanduser(work_dir),
|
||||
# 对齐 L9:POSIX 下 shell 成为会话首进程,kill 按进程组
|
||||
# 一次覆盖整棵进程树(含 & 后台孙进程)
|
||||
start_new_session=(os.name == "posix"),
|
||||
)
|
||||
except FileNotFoundError as e:
|
||||
return {"ok": False, "error": f"命令未找到: {e}"}
|
||||
except PermissionError as e:
|
||||
return {"ok": False, "error": f"权限不足: {e}"}
|
||||
except Exception as e:
|
||||
return {"ok": False, "error": f"启动失败: {e}"}
|
||||
|
||||
pgid = None
|
||||
if os.name == "posix":
|
||||
try:
|
||||
pgid = os.getpgid(popen.pid)
|
||||
except Exception:
|
||||
pgid = None
|
||||
|
||||
entry = _ProcEntry(
|
||||
handle, name, command.strip(), popen, pgid,
|
||||
detach=bool(detach), max_lifetime=self.max_lifetime,
|
||||
buffer_bytes=self.buffer_per_process, owner=owner,
|
||||
)
|
||||
# 可选全档镜像(失败降级为仅环形缓冲,不阻断启动)
|
||||
log_path = None
|
||||
if log_file:
|
||||
try:
|
||||
log_dir = constants.WORKSPACE_DIR
|
||||
os.makedirs(log_dir, exist_ok=True)
|
||||
log_path = os.path.join(log_dir, f"bg_{handle}.log")
|
||||
entry.log_fh = open(log_path, "a", encoding="utf-8",
|
||||
errors="replace", buffering=1)
|
||||
except Exception as e:
|
||||
entry.log_fh = None
|
||||
log_path = None
|
||||
self.logger.warning(f"输出镜像文件开启失败,仅保留环形缓冲: {e}")
|
||||
with self._lock:
|
||||
# start() 竞态兜底:期间并发 shutdown 则立即回收刚拉起的进程
|
||||
if self._shutdown_done:
|
||||
self._close_stdin(entry)
|
||||
self._terminate_tree(entry, force=True)
|
||||
try:
|
||||
popen.wait(timeout=2)
|
||||
except Exception:
|
||||
pass
|
||||
return {"ok": False, "error": "管理器已关闭"}
|
||||
self._procs[handle] = entry
|
||||
|
||||
entry.reader = threading.Thread(
|
||||
target=self._reader, args=(entry,), daemon=True,
|
||||
name=f"bg-proc-{handle}-reader")
|
||||
entry.reader.start()
|
||||
self.logger.info(f"后台进程已启动: #{handle} pid={popen.pid} "
|
||||
f"detach={entry.detach} cmd={entry.command[:60]}")
|
||||
return {"ok": True, "handle": handle, "pid": popen.pid, "name": name,
|
||||
"detach": entry.detach, "log_file": log_path}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 输出读取
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_output(self, handle, tail: int = None, since: int = None,
|
||||
max_bytes: int = 65536) -> Dict:
|
||||
"""拉取后台进程输出。
|
||||
|
||||
- tail=N:取最近 N 行(缺省 50)
|
||||
- since=S:取行序号 > S 的全部缓冲行(增量续读,游标为上次
|
||||
返回的 next_cursor;被环形淘汰的早期行以 evicted 提示)
|
||||
- max_bytes:本次拉取字节上限(缺省 64KB,硬顶 1MB)
|
||||
"""
|
||||
entry = self._lookup(handle)
|
||||
if entry is None:
|
||||
return {"ok": False, "error": f"句柄 #{handle} 不存在"}
|
||||
|
||||
try:
|
||||
max_bytes = int(max_bytes or 65536)
|
||||
except (TypeError, ValueError):
|
||||
max_bytes = 65536
|
||||
max_bytes = max(1024, min(max_bytes, MAX_FETCH_BYTES))
|
||||
|
||||
with entry.lock:
|
||||
snapshot = list(entry.lines)
|
||||
evicted = entry.evicted
|
||||
next_seq = entry.next_seq
|
||||
|
||||
selected = []
|
||||
evicted_note = 0
|
||||
if since is not None:
|
||||
try:
|
||||
since = int(since)
|
||||
except (TypeError, ValueError):
|
||||
return {"ok": False, "error": "since 游标必须是整数"}
|
||||
selected = [(s, l) for (s, l, _) in snapshot if s > since]
|
||||
# 请求范围早于环形缓冲起点:被淘汰部分如实告知
|
||||
oldest = snapshot[0][0] if snapshot else next_seq
|
||||
if since + 1 < oldest:
|
||||
evicted_note = oldest - since - 1
|
||||
else:
|
||||
try:
|
||||
tail = int(tail) if tail is not None else 50
|
||||
except (TypeError, ValueError):
|
||||
tail = 50
|
||||
tail = max(1, min(tail, 100000))
|
||||
selected = [(s, l) for (s, l, _) in snapshot[-tail:]]
|
||||
|
||||
# 字节上限截断(保留最新行优先:从尾部回溯收集)
|
||||
kept = []
|
||||
total = 0
|
||||
truncated = False
|
||||
for s, line in reversed(selected):
|
||||
nb = len(line.encode("utf-8", errors="replace"))
|
||||
if total + nb > max_bytes:
|
||||
truncated = True
|
||||
break
|
||||
kept.append((s, line))
|
||||
total += nb
|
||||
kept.reverse()
|
||||
|
||||
output = "".join(line for _, line in kept)
|
||||
result = {
|
||||
"ok": True,
|
||||
"handle": entry.handle,
|
||||
"name": entry.name,
|
||||
"status": entry.status,
|
||||
"exit_code": entry.exit_code,
|
||||
"output": output,
|
||||
"lines": len(kept),
|
||||
"next_cursor": kept[-1][0] if kept else max(since or 0, next_seq - 1),
|
||||
"truncated": truncated,
|
||||
}
|
||||
if evicted_note > 0 or (since is None and evicted > 0):
|
||||
result["evicted"] = evicted_note if since is not None else evicted
|
||||
return result
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 交互写入(stdin 管道 + watermark 对齐)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def write(self, handle, text: str, append_newline: bool = True,
|
||||
wait_lines: int = 0, wait_timeout: float = 5.0,
|
||||
write_timeout: float = 5.0, owner: str = None) -> Dict:
|
||||
"""向后台进程 stdin 写入一行(交互式问答)。
|
||||
|
||||
返回 {ok, handle, watermark, sent_bytes[, wait]}:
|
||||
- watermark:本条输入的 [stdin] 回显行序号;bg_output(since=
|
||||
watermark) 得到的必然且仅是本条输入引发的响应(写读对齐)
|
||||
- wait_lines>0 时同步等待 N 行响应(wait 子 dict 含 output/
|
||||
timed_out),expect 式一步问答
|
||||
失败返回 {ok: False, error}(进程已退出/写入超时/超长等)。
|
||||
|
||||
T22 属主隔离:条目与调用方都携带 owner 且不一致时拒绝——他人
|
||||
stdin 不容注入。任一侧未标识(None,如测试/管理面直调)不隔离。"""
|
||||
entry = self._lookup(handle)
|
||||
if entry is None:
|
||||
return {"ok": False, "error": f"句柄 #{handle} 不存在"}
|
||||
denial = self._check_owner(entry, owner, "写入")
|
||||
if denial:
|
||||
return denial
|
||||
if not isinstance(text, str) or text == "":
|
||||
return {"ok": False, "error": "缺少或空的 text"}
|
||||
data = text if (not append_newline or text.endswith("\n")) else text + "\n"
|
||||
payload = data.encode("utf-8")
|
||||
if len(payload) > self.max_write_bytes:
|
||||
return {"ok": False,
|
||||
"error": f"单次写入超过 {self.max_write_bytes} 字节上限(超长内容建议落文件后让进程从文件读)"}
|
||||
|
||||
with self._lock:
|
||||
self._update_status_locked(entry)
|
||||
if entry.status != "running":
|
||||
return {"ok": False, "error": f"进程已结束({entry.status}),无法写入"}
|
||||
|
||||
# watermark:取号与回显行记账必须在同一临界区,保证 watermark 等于
|
||||
# 回显行真实序号,bg_output(since=watermark) 正确对齐本条响应(T8)
|
||||
echo = "[stdin] " + text.replace("\n", "⏎")
|
||||
with entry.lock:
|
||||
watermark = entry.next_seq
|
||||
# 回显行进环形缓冲 + log 镜像(多行输入折叠为单行展示)
|
||||
self._record_line(entry, echo)
|
||||
|
||||
# 写入放 watchdog 线程:进程不读 stdin 且管道缓冲(~64KB)写满时
|
||||
# write() 会阻塞,不能拖死工具线程;超时后 daemon 线程由系统兜底
|
||||
error_box = {}
|
||||
done = threading.Event()
|
||||
|
||||
def _do_write():
|
||||
try:
|
||||
with entry.stdin_lock:
|
||||
stdin = entry.popen.stdin
|
||||
if stdin is None or stdin.closed:
|
||||
raise BrokenPipeError("stdin 已关闭")
|
||||
stdin.write(data)
|
||||
stdin.flush()
|
||||
except Exception as e:
|
||||
error_box["e"] = e
|
||||
finally:
|
||||
done.set()
|
||||
|
||||
t = threading.Thread(target=_do_write, daemon=True,
|
||||
name=f"bg-proc-{entry.handle}-writer")
|
||||
t.start()
|
||||
try:
|
||||
write_timeout = max(0.1, min(60.0, float(write_timeout)))
|
||||
except (TypeError, ValueError):
|
||||
write_timeout = 5.0
|
||||
if not done.wait(timeout=write_timeout):
|
||||
return {"ok": False,
|
||||
"error": f"写入超时(>{write_timeout}s):进程可能不读取 stdin,管道缓冲已满"}
|
||||
if error_box:
|
||||
with self._lock:
|
||||
self._update_status_locked(entry)
|
||||
return {"ok": False,
|
||||
"error": f"写入失败: {error_box['e']}"}
|
||||
|
||||
result = {"ok": True, "handle": entry.handle, "watermark": watermark,
|
||||
"sent_bytes": len(payload)}
|
||||
if wait_lines and int(wait_lines) > 0:
|
||||
result["wait"] = self._wait_lines(entry, watermark,
|
||||
int(wait_lines), wait_timeout)
|
||||
return result
|
||||
|
||||
def _wait_lines(self, entry: _ProcEntry, watermark: int, count: int,
|
||||
timeout: float) -> Dict:
|
||||
"""等待 watermark 之后出现 count 行新输出(进程退出且 reader 排空
|
||||
即提前返回,如实标记 timed_out)。"""
|
||||
try:
|
||||
timeout = max(0.1, min(60.0, float(timeout)))
|
||||
except (TypeError, ValueError):
|
||||
timeout = 5.0
|
||||
deadline = time.time() + timeout
|
||||
|
||||
def _snap():
|
||||
with entry.lock:
|
||||
return [(s, l) for (s, l, _) in entry.lines if s > watermark]
|
||||
|
||||
lines = _snap()
|
||||
timed_out = False
|
||||
while len(lines) < count:
|
||||
if entry.popen.poll() is not None and entry.reader_done.is_set():
|
||||
time.sleep(0.1) # 排空 reader 尾部在途行后取终值
|
||||
lines = _snap()
|
||||
break
|
||||
if time.time() > deadline:
|
||||
timed_out = True
|
||||
break
|
||||
time.sleep(0.05)
|
||||
lines = _snap()
|
||||
|
||||
# 字节上限(复用单次拉取硬顶)
|
||||
kept, total, truncated = [], 0, False
|
||||
for s, line in lines:
|
||||
nb = len(line.encode("utf-8", errors="replace"))
|
||||
if total + nb > MAX_FETCH_BYTES:
|
||||
truncated = True
|
||||
break
|
||||
kept.append((s, line))
|
||||
total += nb
|
||||
out = {
|
||||
"lines": len(kept),
|
||||
"output": "".join(l for _, l in kept),
|
||||
"timed_out": timed_out,
|
||||
"next_cursor": kept[-1][0] if kept else watermark,
|
||||
}
|
||||
if truncated:
|
||||
out["truncated"] = True
|
||||
return out
|
||||
|
||||
def _close_stdin(self, entry: _ProcEntry):
|
||||
"""关闭 stdin 管道(行协议进程见 EOF 常自行退出;幂等)"""
|
||||
try:
|
||||
with entry.stdin_lock:
|
||||
stdin = entry.popen.stdin
|
||||
if stdin is not None and not stdin.closed:
|
||||
stdin.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 列举 / 终止 / 收尾
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def list_procs(self) -> List[Dict]:
|
||||
"""列出全部条目(运行中在前,其余按句柄序)"""
|
||||
with self._lock:
|
||||
self._sweep_locked()
|
||||
# 惰性收割:列举即刷新全部条目状态(此前仅 kill/get_output 路径更新)
|
||||
for entry in self._procs.values():
|
||||
self._update_status_locked(entry)
|
||||
entries = sorted(self._procs.values(),
|
||||
key=lambda e: (e.status != "running", e.handle))
|
||||
return [self._entry_info(e) for e in entries]
|
||||
|
||||
def kill(self, handle, force: bool = False, grace: float = 3.0,
|
||||
owner: str = None) -> Dict:
|
||||
"""终止后台进程。SIGTERM → 宽限轮询 → SIGKILL;force=True 直接 SIGKILL。
|
||||
对已退出条目幂等,返回当前状态。
|
||||
|
||||
T22 属主隔离:条目与调用方都携带 owner 且不一致时拒绝——防止
|
||||
提示注入的会话/任务跨属主杀掉他人的 dev server。shutdown/TTL
|
||||
清理属管理器特权,不走本校验。"""
|
||||
entry = self._lookup(handle)
|
||||
if entry is None:
|
||||
return {"ok": False, "error": f"句柄 #{handle} 不存在"}
|
||||
denial = self._check_owner(entry, owner, "终止")
|
||||
if denial:
|
||||
return denial
|
||||
|
||||
with self._lock:
|
||||
self._update_status_locked(entry)
|
||||
if entry.status != "running":
|
||||
return {"ok": True, "handle": entry.handle, "status": entry.status,
|
||||
"exit_code": entry.exit_code, "already_finished": True}
|
||||
|
||||
try:
|
||||
grace = min(30.0, max(0.1, float(grace)))
|
||||
except (TypeError, ValueError):
|
||||
grace = 3.0
|
||||
|
||||
if force:
|
||||
self._terminate_tree(entry, force=True)
|
||||
else:
|
||||
# 三级优雅终止:关 stdin(行协议进程见 EOF 常自行退出)→
|
||||
# SIGTERM 宽限 → SIGKILL 兜底
|
||||
self._close_stdin(entry)
|
||||
deadline = time.time() + grace * 0.4
|
||||
while time.time() < deadline and entry.popen.poll() is None:
|
||||
time.sleep(0.05)
|
||||
if entry.popen.poll() is None:
|
||||
self._terminate_tree(entry)
|
||||
deadline = time.time() + grace * 0.6
|
||||
while time.time() < deadline and entry.popen.poll() is None:
|
||||
time.sleep(0.05)
|
||||
if entry.popen.poll() is None:
|
||||
self._terminate_tree(entry, force=True)
|
||||
try:
|
||||
entry.popen.wait(timeout=2)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
with self._lock:
|
||||
self._update_status_locked(entry)
|
||||
if entry.status == "running":
|
||||
# 进程仍未死(极端:不可杀的 D 状态),如实报告
|
||||
return {"ok": False, "handle": entry.handle,
|
||||
"error": "进程未能终止(可能处于不可中断状态)"}
|
||||
entry.status = "killed"
|
||||
return {"ok": True, "handle": entry.handle, "status": entry.status,
|
||||
"exit_code": entry.exit_code}
|
||||
|
||||
def shutdown(self) -> Dict:
|
||||
"""退出收尾:终止全部非 detach 运行中进程(SIGTERM→SIGKILL),
|
||||
清空注册表;detach 进程保留运行、仅注销登记。幂等。"""
|
||||
with self._lock:
|
||||
if self._shutdown_done:
|
||||
return {"killed": 0, "detached": 0}
|
||||
self._shutdown_done = True
|
||||
entries = list(self._procs.values())
|
||||
self._procs.clear()
|
||||
|
||||
killed = detached = 0
|
||||
for entry in entries:
|
||||
self._update_status(entry)
|
||||
if entry.status != "running":
|
||||
continue
|
||||
if entry.detach:
|
||||
detached += 1
|
||||
self.logger.info(f"后台进程 #{entry.handle}(detach)保留运行: "
|
||||
f"pid={entry.popen.pid}")
|
||||
continue
|
||||
# 优雅收尾:先关 stdin(行协议进程见 EOF 常自行退出)再走信号
|
||||
self._close_stdin(entry)
|
||||
try:
|
||||
entry.popen.wait(timeout=1)
|
||||
except Exception:
|
||||
pass
|
||||
if entry.popen.poll() is None:
|
||||
self._terminate_tree(entry)
|
||||
try:
|
||||
entry.popen.wait(timeout=3)
|
||||
except Exception:
|
||||
self._terminate_tree(entry, force=True)
|
||||
try:
|
||||
entry.popen.wait(timeout=2)
|
||||
except Exception:
|
||||
pass
|
||||
entry.status = "killed"
|
||||
killed += 1
|
||||
if killed or detached:
|
||||
self.logger.info(f"后台进程收尾: 终止 {killed} 个, detach 保留 {detached} 个")
|
||||
return {"killed": killed, "detached": detached}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 内部实现
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _check_owner(self, entry: _ProcEntry, caller_owner: Optional[str],
|
||||
action: str) -> Optional[Dict]:
|
||||
"""T22:属主校验——条目与调用方都携带 owner 且不一致时返回拒绝
|
||||
dict(写/杀权限隔离);任一侧未标识(None,如测试/管理面直调)
|
||||
不隔离,保持向后兼容。观察面(list/get_output)不经过本校验。"""
|
||||
if (entry.owner is not None and caller_owner is not None
|
||||
and entry.owner != caller_owner):
|
||||
self.logger.warning(
|
||||
f"属主校验拒绝:owner={caller_owner} 尝试{action} "
|
||||
f"#{entry.handle}(属主 {entry.owner})")
|
||||
return {"ok": False,
|
||||
"error": f"进程 #{entry.handle} 属于其他会话(属主不匹配),无权{action}"}
|
||||
return None
|
||||
|
||||
def _lookup(self, handle) -> Optional[_ProcEntry]:
|
||||
try:
|
||||
handle = int(handle)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
with self._lock:
|
||||
entry = self._procs.get(handle)
|
||||
if entry is not None:
|
||||
self._update_status(entry)
|
||||
return entry
|
||||
|
||||
def _update_status(self, entry: _ProcEntry):
|
||||
with self._lock:
|
||||
self._update_status_locked(entry)
|
||||
|
||||
def _update_status_locked(self, entry: _ProcEntry):
|
||||
"""惰性收割:poll() 即 waitpid(WNOHANG),自然退出时记录状态。"""
|
||||
if entry.status != "running":
|
||||
return
|
||||
rc = entry.popen.poll()
|
||||
if rc is not None:
|
||||
entry.exit_code = rc
|
||||
entry.end_time = time.time()
|
||||
entry.status = "exited"
|
||||
|
||||
def _sweep_locked(self):
|
||||
"""TTL 惰性清理:超期运行中进程标记 expired 并终止(max_lifetime<=0 不启用)。
|
||||
调用方须持有 self._lock。"""
|
||||
if self.max_lifetime <= 0:
|
||||
return
|
||||
now = time.time()
|
||||
for entry in list(self._procs.values()):
|
||||
if entry.status == "running" and now - entry.started_at > self.max_lifetime:
|
||||
self._terminate_tree(entry)
|
||||
entry.status = "expired"
|
||||
entry.end_time = now
|
||||
entry.exit_code = entry.popen.poll()
|
||||
self.logger.info(f"后台进程 #{entry.handle} TTL 到期已终止")
|
||||
|
||||
def _terminate_tree(self, entry: _ProcEntry, force: bool = False):
|
||||
"""按进程组终止整棵进程树(对齐 executors L9/T11 语义与 Windows 回退)"""
|
||||
proc = entry.popen
|
||||
try:
|
||||
if os.name == "posix":
|
||||
sig = signal.SIGKILL if force else signal.SIGTERM
|
||||
pgid = entry.pgid
|
||||
if pgid is None:
|
||||
pgid = os.getpgid(proc.pid)
|
||||
os.killpg(pgid, sig)
|
||||
else:
|
||||
# T11(审计 2026-08-21):Windows 无进程组,单进程 kill 收割
|
||||
# 不到 spawn 的孙进程(dev server/训练任务成孤儿继续占用端口
|
||||
# /文件锁)。taskkill /T /F 递归终止整棵树,失败兜底单进程 kill
|
||||
try:
|
||||
subprocess.run(
|
||||
["taskkill", "/PID", str(proc.pid), "/T", "/F"],
|
||||
capture_output=True, timeout=10)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
if proc.poll() is None:
|
||||
proc.kill()
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
try:
|
||||
proc.kill()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _reader(self, entry: _ProcEntry):
|
||||
"""后台 reader 线程(对齐 executors H2 模式):阻塞 readline 入环形
|
||||
缓冲(+ 可选 log 镜像);EOF/异常即 poll() 收割防僵尸,并显式关闭
|
||||
管道与镜像文件句柄。"""
|
||||
try:
|
||||
for line in iter(entry.popen.stdout.readline, ""):
|
||||
self._record_line(entry, line)
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
entry.reader_done.set()
|
||||
try:
|
||||
entry.popen.poll() # 立即收割退出子进程,防僵尸悬挂
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
entry.popen.stdout.close()
|
||||
except Exception:
|
||||
pass
|
||||
if entry.log_fh is not None:
|
||||
try:
|
||||
entry.log_fh.flush()
|
||||
entry.log_fh.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _record_line(self, entry: _ProcEntry, line: str):
|
||||
"""统一记录:环形缓冲 + 可选 log 全档镜像(stdin 回显同样入档)"""
|
||||
if entry.log_fh is not None:
|
||||
try:
|
||||
entry.log_fh.write(line)
|
||||
except Exception:
|
||||
pass
|
||||
self._append_line(entry, line)
|
||||
|
||||
@staticmethod
|
||||
def _append_line(entry: _ProcEntry, line: str):
|
||||
# 单行截断:环形缓冲内存预算不被超长单行击穿
|
||||
nbytes = len(line.encode("utf-8", errors="replace"))
|
||||
if nbytes > MAX_LINE_BYTES:
|
||||
allowed = line.encode("utf-8", errors="replace")[:MAX_LINE_BYTES]
|
||||
line = allowed.decode("utf-8", errors="ignore") + "…[单行截断]\n"
|
||||
nbytes = len(line.encode("utf-8", errors="replace"))
|
||||
with entry.lock:
|
||||
entry.lines.append((entry.next_seq, line, nbytes))
|
||||
entry.next_seq += 1
|
||||
entry.buffer_bytes += nbytes
|
||||
while entry.buffer_bytes > entry.max_buffer and len(entry.lines) > 1:
|
||||
_, old, old_nb = entry.lines.popleft()
|
||||
entry.buffer_bytes -= old_nb
|
||||
entry.evicted += 1
|
||||
|
||||
@staticmethod
|
||||
def _entry_info(entry: _ProcEntry) -> Dict:
|
||||
now = time.time()
|
||||
return {
|
||||
"handle": entry.handle,
|
||||
"name": entry.name,
|
||||
"pid": entry.popen.pid,
|
||||
"status": entry.status,
|
||||
"exit_code": entry.exit_code,
|
||||
"detach": entry.detach,
|
||||
"started_at": entry.started_at,
|
||||
"runtime_sec": round((entry.end_time or now) - entry.started_at, 1),
|
||||
"buffered_lines": len(entry.lines),
|
||||
"evicted_lines": entry.evicted,
|
||||
"command": entry.command,
|
||||
"owner": entry.owner, # T22:属主标识(观察面可见,写/杀需匹配)
|
||||
}
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# 模块级单例(工具层共享;测试经 reset_process_manager() 重置)
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
_manager: Optional[ProcessManager] = None
|
||||
_manager_lock = threading.Lock()
|
||||
_atexit_registered = False
|
||||
|
||||
|
||||
def get_process_manager(logger: Optional[Logger] = None) -> ProcessManager:
|
||||
"""获取全局 ProcessManager 单例(懒加载)。
|
||||
|
||||
后台进程注册表必须跨会话/跨子代理共享——任何 worker 子代理启动的
|
||||
进程,后续任意子代理都能 bg_list/bg_output/bg_kill 管理。"""
|
||||
global _manager, _atexit_registered
|
||||
with _manager_lock:
|
||||
if _manager is None:
|
||||
_manager = ProcessManager(logger)
|
||||
if not _atexit_registered:
|
||||
# 退出语义:aibox 进程退出(含 --once sys.exit / Ctrl+C)时
|
||||
# 非 detach 后台进程统一终止;detach 进程保留运行仅注销登记。
|
||||
# SIGKILL/SIGTERM 硬杀无法覆盖(与既有线程池等资源同级)。
|
||||
atexit.register(_shutdown_atexit)
|
||||
_atexit_registered = True
|
||||
return _manager
|
||||
|
||||
|
||||
def reset_process_manager(manager: Optional[ProcessManager] = None) -> None:
|
||||
"""替换/清空全局单例(测试隔离用)。传入 None 仅解除引用,不触发收尾。"""
|
||||
global _manager
|
||||
with _manager_lock:
|
||||
_manager = manager
|
||||
|
||||
|
||||
def _shutdown_atexit():
|
||||
m = _manager
|
||||
if m is not None:
|
||||
try:
|
||||
m.shutdown()
|
||||
except Exception:
|
||||
pass
|
||||
@@ -18,10 +18,12 @@ subsystems.prompt_manager - 提示词管理器
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import threading
|
||||
from datetime import datetime
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from common import ConfigManager
|
||||
from common.atomic import atomic_write_text
|
||||
|
||||
|
||||
class PromptManager:
|
||||
@@ -35,6 +37,7 @@ class PromptManager:
|
||||
"""
|
||||
|
||||
_instance = None
|
||||
_instance_lock = threading.Lock() # I 批:并发首例化竞态防护
|
||||
_prompts = []
|
||||
_prompts_dir = None
|
||||
|
||||
@@ -43,27 +46,43 @@ class PromptManager:
|
||||
STATUS_FILE = "100_STATUS.md"
|
||||
SOUL_FILE = "99_SOUL.md" # 智能体灵魂文件 - 所有角色定位内容
|
||||
|
||||
# N9(审计 2026-08-18):Agent 可控内容直写 00/99/100 文件并作为后续
|
||||
# 会话系统提示词加载(持久化注入面)。内容长度上限约束注入规模,
|
||||
# 超长内容拒绝写入(工具层有更友好的提示)。
|
||||
MAX_CONTENT_LENGTH = 10000
|
||||
|
||||
# 基础技能文件名模式:N_xxx.md(N 为数字,如 1_memory_search.md、10_task_handling.md)
|
||||
_SKILL_FILE_RE = re.compile(r'^(\d+)_')
|
||||
|
||||
def __new__(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
with cls._instance_lock:
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
return cls._instance
|
||||
|
||||
def __init__(self):
|
||||
if not self._prompts:
|
||||
self._load_prompts()
|
||||
# J6(审计 2026-08-19):__new__ 的锁只防双实例,多线程首例并发
|
||||
# 同时进入 __init__ 会使 _load_prompts 双重执行,_prompts(类属性)
|
||||
# 整份翻倍——初始化检查与加载收进同一把锁
|
||||
with self._instance_lock:
|
||||
if not self._prompts:
|
||||
self._load_prompts()
|
||||
|
||||
def _get_prompts_dir(self) -> str:
|
||||
config = ConfigManager()
|
||||
prompts_dir = config.get("system.prompts_dir", "~/.aibox/prompts")
|
||||
return os.path.expanduser(prompts_dir)
|
||||
|
||||
def _load_prompts(self):
|
||||
def _load_prompts(self, target: Optional[List[Dict]] = None):
|
||||
"""加载提示词到 target 列表(缺省 self._prompts)。
|
||||
|
||||
M6:reload 传入新列表本地构建,完成后由 reload 原子替换——
|
||||
此前先清空 self._prompts 再逐文件重读,并发读者(其他连接构建
|
||||
system 消息)会读到空/半截提示词列表。
|
||||
"""
|
||||
prompts = target if target is not None else self._prompts
|
||||
self._prompts_dir = self._get_prompts_dir()
|
||||
os.makedirs(self._prompts_dir, exist_ok=True)
|
||||
self._prompts = []
|
||||
|
||||
# 确保核心文件存在
|
||||
self._ensure_identity_file()
|
||||
@@ -86,17 +105,14 @@ class PromptManager:
|
||||
with open(filepath, 'r', encoding='utf-8') as f:
|
||||
content = f.read().strip()
|
||||
if content:
|
||||
self._prompts.append({
|
||||
prompts.append({
|
||||
'file': filename,
|
||||
'content': content
|
||||
})
|
||||
except Exception as e:
|
||||
print(f"⚠️ 读取提示词文件 {filename} 失败: {e}")
|
||||
|
||||
essential_files = [self.IDENTITY_FILE, self.STATUS_FILE, self.SOUL_FILE]
|
||||
other_files = [f for f in prompt_files if f not in essential_files] if prompt_files else []
|
||||
|
||||
if not self._prompts:
|
||||
if not prompts:
|
||||
self._create_default_prompts()
|
||||
|
||||
def _get_template_dir(self) -> str:
|
||||
@@ -148,6 +164,9 @@ class PromptManager:
|
||||
|
||||
def update_identity(self, content: str) -> bool:
|
||||
"""更新用户身份信息"""
|
||||
if len(content or "") > self.MAX_CONTENT_LENGTH:
|
||||
print(f"⚠️ 身份内容超长({len(content)} > {self.MAX_CONTENT_LENGTH} 字符),拒绝写入")
|
||||
return False
|
||||
identity_path = os.path.join(self._prompts_dir, self.IDENTITY_FILE)
|
||||
try:
|
||||
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
@@ -158,8 +177,8 @@ class PromptManager:
|
||||
【更新记录】
|
||||
- 上次更新:{timestamp}
|
||||
"""
|
||||
with open(identity_path, 'w', encoding='utf-8') as f:
|
||||
f.write(updated_content)
|
||||
# J5:原子写,崩溃不留半截提示词文件
|
||||
atomic_write_text(identity_path, updated_content)
|
||||
self.reload()
|
||||
return True
|
||||
except Exception as e:
|
||||
@@ -168,6 +187,9 @@ class PromptManager:
|
||||
|
||||
def update_status(self, content: str) -> bool:
|
||||
"""更新智能体状态"""
|
||||
if len(content or "") > self.MAX_CONTENT_LENGTH:
|
||||
print(f"⚠️ 状态内容超长({len(content)} > {self.MAX_CONTENT_LENGTH} 字符),拒绝写入")
|
||||
return False
|
||||
status_path = os.path.join(self._prompts_dir, self.STATUS_FILE)
|
||||
try:
|
||||
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
@@ -178,8 +200,8 @@ class PromptManager:
|
||||
【更新记录】
|
||||
- 上次更新:{timestamp}
|
||||
"""
|
||||
with open(status_path, 'w', encoding='utf-8') as f:
|
||||
f.write(updated_content)
|
||||
# J5:原子写,崩溃不留半截提示词文件
|
||||
atomic_write_text(status_path, updated_content)
|
||||
self.reload()
|
||||
return True
|
||||
except Exception as e:
|
||||
@@ -188,6 +210,9 @@ class PromptManager:
|
||||
|
||||
def update_soul(self, content: str) -> bool:
|
||||
"""更新智能体灵魂定义 - 所有角色定位内容"""
|
||||
if len(content or "") > self.MAX_CONTENT_LENGTH:
|
||||
print(f"⚠️ 灵魂内容超长({len(content)} > {self.MAX_CONTENT_LENGTH} 字符),拒绝写入")
|
||||
return False
|
||||
soul_path = os.path.join(self._prompts_dir, self.SOUL_FILE)
|
||||
try:
|
||||
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
@@ -198,8 +223,8 @@ class PromptManager:
|
||||
【更新记录】
|
||||
- 上次更新:{timestamp}
|
||||
"""
|
||||
with open(soul_path, 'w', encoding='utf-8') as f:
|
||||
f.write(updated_content)
|
||||
# J5:原子写,崩溃不留半截提示词文件
|
||||
atomic_write_text(soul_path, updated_content)
|
||||
self.reload()
|
||||
return True
|
||||
except Exception as e:
|
||||
@@ -240,10 +265,6 @@ class PromptManager:
|
||||
return """我是DeepSeek,一个智能AI助手。
|
||||
请创建 99_SOUL.md 文件来定义我的完整身份、角色、性格和能力边界。"""
|
||||
|
||||
def get_skill_prompts(self) -> List[Dict]:
|
||||
"""获取所有基础技能提示词(以1_开头的文件)"""
|
||||
return [p for p in self._prompts if p['file'].startswith(self.SKILL_PREFIX)]
|
||||
|
||||
def get_soul_prompt(self) -> Optional[Dict]:
|
||||
"""获取灵魂提示词(99_SOUL.md)"""
|
||||
for p in self._prompts:
|
||||
@@ -296,9 +317,10 @@ class PromptManager:
|
||||
return [p for p in self._prompts if self._skill_sort_key(p['file']) is not None]
|
||||
|
||||
def reload(self):
|
||||
"""重新加载所有提示词"""
|
||||
self._prompts = []
|
||||
self._load_prompts()
|
||||
"""重新加载所有提示词(M6:构建新列表后原子替换,并发读者始终读到完整列表)"""
|
||||
new_prompts: List[Dict] = []
|
||||
self._load_prompts(new_prompts)
|
||||
self._prompts = new_prompts
|
||||
return len(self._prompts)
|
||||
|
||||
def get_prompt_count(self) -> int:
|
||||
|
||||
@@ -27,6 +27,9 @@ class SkillManager:
|
||||
2. Markdown 文件 (.md):第一行为简介(# 标题或普通文本),其余为技能内容
|
||||
"""
|
||||
|
||||
# J16(审计 2026-08-19):--help 输出上限——内容会进系统提示词
|
||||
MAX_HELP_LENGTH = 2000
|
||||
|
||||
def __init__(self, logger: Optional[Logger] = None):
|
||||
self.logger = logger or Logger("SkillManager")
|
||||
# name -> {path, file, summary, content, is_markdown}
|
||||
@@ -38,7 +41,11 @@ class SkillManager:
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _run_help(self, skill_path: str) -> str:
|
||||
"""执行 skill_path --help 并返回完整输出(stdout 优先,否则 stderr)。"""
|
||||
"""执行 skill_path --help 并返回完整输出(stdout 优先,否则 stderr)。
|
||||
|
||||
J16(审计 2026-08-19):输出截断到 MAX_HELP_LENGTH——技能 --help
|
||||
输出会进系统提示词,异常/恶意脚本可撑爆上下文。
|
||||
"""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[skill_path, "--help"],
|
||||
@@ -50,6 +57,8 @@ class SkillManager:
|
||||
output = result.stdout.strip()
|
||||
if not output:
|
||||
output = result.stderr.strip()
|
||||
if len(output) > self.MAX_HELP_LENGTH:
|
||||
output = output[:self.MAX_HELP_LENGTH] + "\n…(已截断)"
|
||||
return output
|
||||
except Exception as e:
|
||||
self.logger.warning(f"运行 {skill_path} --help 失败: {e}")
|
||||
@@ -114,12 +123,18 @@ class SkillManager:
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _load_skills(self):
|
||||
"""扫描技能目录,加载所有技能文件(可执行文件或 Markdown 文件)。"""
|
||||
"""扫描技能目录,加载所有技能文件(可执行文件或 Markdown 文件)。
|
||||
|
||||
T26(审计 2026-08-21):本地构建 new_skills 完成后一次性原子替换,
|
||||
对齐 PromptManager M6——此前先 self.skills = {} 清空再逐个重建,
|
||||
reload 期间(每个可执行技能 --help 最多 10s)并发读者会看到空/半截
|
||||
技能表(技能"消失"再"逐渐出现")。"""
|
||||
skills_dir = constants.SKILLS_DIR
|
||||
self.skills = {}
|
||||
new_skills = {}
|
||||
|
||||
if not os.path.exists(skills_dir):
|
||||
self.logger.info(f"技能目录不存在: {skills_dir}")
|
||||
self.skills = new_skills
|
||||
return
|
||||
|
||||
skill_files = [
|
||||
@@ -129,16 +144,30 @@ class SkillManager:
|
||||
and not f.startswith('.')
|
||||
]
|
||||
|
||||
for skill_file in skill_files:
|
||||
# J16(审计 2026-08-19):同名键冲突(foo 可执行文件与 foo.md)此前
|
||||
# 由 listdir 顺序决定谁覆盖谁,行为不确定。改为可执行优先 + 确定性:
|
||||
# 先加载可执行技能,Markdown 遇到已占用的同名键时跳过并告警
|
||||
executable_files = [
|
||||
f for f in skill_files
|
||||
if not self._is_markdown_file(f)
|
||||
]
|
||||
markdown_files = [f for f in skill_files if self._is_markdown_file(f)]
|
||||
|
||||
for skill_file in executable_files + markdown_files:
|
||||
skill_path = os.path.join(skills_dir, skill_file)
|
||||
skill_name = os.path.splitext(skill_file)[0]
|
||||
is_markdown = self._is_markdown_file(skill_file)
|
||||
|
||||
if is_markdown:
|
||||
if skill_name in new_skills:
|
||||
self.logger.warning(
|
||||
f"Markdown 技能 '{skill_file}' 与同名可执行技能冲突,已跳过"
|
||||
f"(可执行技能优先)")
|
||||
continue
|
||||
# Markdown 格式技能:直接读取内容
|
||||
summary, content = self._read_markdown_skill(skill_path)
|
||||
if content:
|
||||
self.skills[skill_name] = {
|
||||
new_skills[skill_name] = {
|
||||
"path": skill_path,
|
||||
"file": skill_file,
|
||||
"summary": summary,
|
||||
@@ -159,7 +188,7 @@ class SkillManager:
|
||||
help_output = self._run_help(skill_path)
|
||||
summary = self._first_nonempty_line(help_output) or f"扩展技能: {skill_name}"
|
||||
|
||||
self.skills[skill_name] = {
|
||||
new_skills[skill_name] = {
|
||||
"path": skill_path,
|
||||
"file": skill_file,
|
||||
"summary": summary,
|
||||
@@ -169,7 +198,8 @@ class SkillManager:
|
||||
}
|
||||
self.logger.info(f"已加载可执行技能: {skill_name} — {summary}")
|
||||
|
||||
self.logger.info(f"共加载 {len(self.skills)} 个技能({len([s for s in self.skills.values() if s['is_markdown']])} 个 Markdown 技能)")
|
||||
self.skills = new_skills
|
||||
self.logger.info(f"共加载 {len(new_skills)} 个技能({len([s for s in new_skills.values() if s['is_markdown']])} 个 Markdown 技能)")
|
||||
|
||||
def reload(self) -> int:
|
||||
"""重新扫描技能目录,返回加载的技能数量。"""
|
||||
|
||||
@@ -0,0 +1,413 @@
|
||||
"""aibox 第四轮审计 Batch H 修复回归测试(2026-08-18,审计报告 §十一)
|
||||
|
||||
覆盖:
|
||||
- H1:子代理共享父取消事件后首轮 think 不再清除父事件(N5 回归)+ 父已取消时跳过启动
|
||||
- M3:调度器任务级超时兜底(AI think 卡死不再停摆调度循环)
|
||||
- M4:get_recent_memories(limit=None) 使用默认上限(不再全表加载)
|
||||
- M5:memo delete 物理删除(此前只标记完成)
|
||||
- M6:PromptManager.reload 原子替换(并发读者不再读到空/半截列表)
|
||||
- L1:read_file 的 sed 参数强转 int(防引号注入)
|
||||
- L2:set_proxy 落盘失败如实报告(不再静默报"已保存")
|
||||
- L3:SSE 非 UTF-8 行不中断流(errors=replace)
|
||||
- L4:persona 路径校验解析符号链接(软链越界拒绝)
|
||||
- L5:1MB 截断不再产生半个字符(�)
|
||||
- L6:compressor 空内容压缩不再除零
|
||||
- L7:InstallSkillTool 指引命令路径引用(防拆词/注入)
|
||||
- L8:ask_user 缺参返回明确错误(不再 TypeError)
|
||||
- L10:web 工具仅允许 http/https(file:// 等拒绝)
|
||||
- L11:registry 日志打码改包含匹配(apiKey/access_token 等变体)
|
||||
- L12:FailoverClient 日志 Bearer 令牌脱敏
|
||||
|
||||
M1(/reload_prompts 移线程池)/M2(技能扫描启动前预热)为 async WebSocket /
|
||||
启动路径,依赖事件循环与真实连接对象,由既有 WebSocket 测试面与启动逻辑回归。
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import queue
|
||||
import sys
|
||||
import threading
|
||||
import types
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
# ---- 环境桩:websockets 未安装时保证 import aibox.py 不失败 ----
|
||||
if "websockets" not in sys.modules:
|
||||
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||
|
||||
import aibox # noqa: E402
|
||||
from aibox import DeepSeekChat, TaskScheduler, run_full_subagent # noqa: E402
|
||||
|
||||
from llm_client import DeepSeekClient, FailoverClient, _sanitize_log_msg # noqa: E402
|
||||
from subsystems.db import MemoryDatabase, MemoDatabase # noqa: E402
|
||||
from subsystems.prompt_manager import PromptManager # noqa: E402
|
||||
from subsystems.conversation_compressor import ConversationCompressor # noqa: E402
|
||||
from common import constants # noqa: E402
|
||||
|
||||
from test_conversation_compressor import FakeLLMClient # noqa: E402
|
||||
from test_once_mode import _make_ai # noqa: E402
|
||||
|
||||
|
||||
def _fake_sub_chat():
|
||||
fake = mock.MagicMock()
|
||||
fake.think_and_respond = mock.MagicMock(return_value="任务完成报告")
|
||||
fake.stats = {}
|
||||
fake.messages = []
|
||||
fake._cancel_event = threading.Event()
|
||||
return fake
|
||||
|
||||
|
||||
# ==================== H1:子代理取消事件清除回归 ====================
|
||||
|
||||
class TestH1SubagentCancelClear:
|
||||
def test_shared_event_not_cleared_on_first_think(self):
|
||||
"""共享父事件后首轮 think 不清除事件——父级已置位的取消信号不被吞掉"""
|
||||
ai = _make_ai()
|
||||
parent_event = threading.Event()
|
||||
parent_event.set()
|
||||
ai._cancel_event = parent_event
|
||||
ai._shared_cancel_event = True
|
||||
with pytest.raises(InterruptedError):
|
||||
ai._think_and_respond_impl("hi")
|
||||
assert parent_event.is_set()
|
||||
|
||||
def test_own_event_still_cleared_on_first_think(self):
|
||||
"""未共享(普通会话)时首轮仍清除自身事件——原语义不回归"""
|
||||
import contextlib
|
||||
ai = _make_ai()
|
||||
ai._cancel_event.set()
|
||||
with contextlib.suppress(Exception):
|
||||
ai._think_and_respond_impl("hi")
|
||||
assert not ai._cancel_event.is_set()
|
||||
|
||||
def test_subagent_not_started_when_parent_already_cancelled(self):
|
||||
"""父事件已置位 → 直接返回取消提示,不启动子代理"""
|
||||
fake = _fake_sub_chat()
|
||||
parent_event = threading.Event()
|
||||
parent_event.set()
|
||||
with mock.patch("aibox.DeepSeekChat", return_value=fake), \
|
||||
mock.patch("aibox.create_llm_client"), \
|
||||
mock.patch("aibox.SubprocessCommandExecutor"):
|
||||
result = run_full_subagent("r", "t", None, cancel_event=parent_event)
|
||||
assert "取消" in result
|
||||
fake.think_and_respond.assert_not_called()
|
||||
|
||||
|
||||
# ==================== M3:调度器任务超时 ====================
|
||||
|
||||
class _BlockingAI:
|
||||
"""think 卡死直到被 cancel 的伪 AI(模拟 LLM 挂起)"""
|
||||
|
||||
def __init__(self):
|
||||
self.cancelled = threading.Event()
|
||||
self.cancel_called = 0
|
||||
|
||||
def cancel(self):
|
||||
self.cancel_called += 1
|
||||
self.cancelled.set()
|
||||
|
||||
def think_and_respond(self, input_text, output_callback=None,
|
||||
is_continuation=False):
|
||||
self.cancelled.wait(timeout=15)
|
||||
raise InterruptedError("用户中断")
|
||||
|
||||
|
||||
class _FakeMemoDB:
|
||||
def get_memo(self, memo_id):
|
||||
return None
|
||||
|
||||
|
||||
class TestM3SchedulerTimeout:
|
||||
def test_stuck_task_times_out_and_scheduler_continues(self):
|
||||
from aibox import ConfigManager
|
||||
loop = asyncio.new_event_loop()
|
||||
try:
|
||||
scheduler = TaskScheduler(loop=loop, logger=None)
|
||||
fake_ai = _BlockingAI()
|
||||
memo = {"id": 7, "title": "卡死任务", "content": "内容", "trigger_count": 0}
|
||||
orig_get = ConfigManager.get
|
||||
|
||||
def fake_get(self, key, default=None):
|
||||
# 只把任务超时压到 0.3s,其余配置读取走原实现(防止
|
||||
# DATETIME_FORMAT 等属性读到错误类型)
|
||||
if key == "scheduler.task_timeout":
|
||||
return 0.3
|
||||
return orig_get(self, key, default)
|
||||
|
||||
with mock.patch.object(ConfigManager, "get", fake_get), \
|
||||
mock.patch.object(scheduler, "_create_task_ai", return_value=fake_ai):
|
||||
result = loop.run_until_complete(
|
||||
scheduler.execute_task_as_conversation(memo, _FakeMemoDB()))
|
||||
assert "超时" in result
|
||||
assert fake_ai.cancel_called >= 1
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
|
||||
# ==================== M4/M5:记忆上限与任务删除 ====================
|
||||
|
||||
class TestM4M5Db:
|
||||
def test_recent_memories_none_capped(self, tmp_path, monkeypatch):
|
||||
"""limit=None 不再全表加载,使用默认上限(对齐 search_memories)"""
|
||||
from test_conversation_compressor import patch_const
|
||||
db = MemoryDatabase(db_path=str(tmp_path / "m.db"))
|
||||
for i in range(7):
|
||||
db.add_memory(topic=f"t{i}", summary=f"s{i}", content="c",
|
||||
start_time="2026-08-18 00:00:00",
|
||||
end_time="2026-08-18 00:00:01", importance=1)
|
||||
patch_const(monkeypatch, "MEMORY_MAX_SEARCH_RESULTS", 5)
|
||||
res = db.get_recent_memories(limit=None)
|
||||
assert len(res) == 5
|
||||
|
||||
def test_delete_memo_physically_removes(self, tmp_path):
|
||||
"""delete_memo 真实删除行(此前 complete_memo 只标记完成)"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
mid = db.add_memo(title="待删任务", content="内容")
|
||||
assert db.get_memo(mid) is not None
|
||||
assert db.delete_memo(mid) is True
|
||||
assert db.get_memo(mid) is None
|
||||
assert db.delete_memo(mid) is False # 已不存在
|
||||
|
||||
def test_memo_tool_delete_uses_physical_delete(self, tmp_path):
|
||||
"""MemoTool delete 走 delete_memo:行消失而非仅标记完成"""
|
||||
from tools import MemoTool
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
mid = db.add_memo(title="重复任务", content="内容", repeat_type="daily")
|
||||
tool = MemoTool(db, None)
|
||||
result = tool.execute(action="delete", memo_id=mid, confirm=True)
|
||||
assert "已删除" in result
|
||||
assert db.get_memo(mid) is None
|
||||
stats = db.get_memo_stats()
|
||||
assert stats["total"] == 0
|
||||
|
||||
|
||||
# ==================== M6:PromptManager 原子 reload ====================
|
||||
|
||||
class TestM6PromptReloadAtomic:
|
||||
def test_reload_builds_new_list_then_swaps(self):
|
||||
"""reload 在本地新列表构建完成后才原子替换——读者始终看到完整列表"""
|
||||
pm = PromptManager()
|
||||
old_list = pm._prompts
|
||||
seen = []
|
||||
|
||||
def fake_load(target):
|
||||
seen.append(target)
|
||||
target.append({"file": "X_test.md", "content": "y"})
|
||||
|
||||
with mock.patch.object(pm, "_load_prompts", side_effect=fake_load):
|
||||
n = pm.reload()
|
||||
assert len(seen) == 1
|
||||
assert seen[0] is not old_list # 构建用的是新列表
|
||||
assert pm._prompts is seen[0] # 完成后才替换
|
||||
assert n == 1
|
||||
|
||||
|
||||
# ==================== L1:read_file 参数强转 ====================
|
||||
|
||||
class TestL1ReadFileCoercion:
|
||||
def test_non_integer_lines_rejected(self, tmp_path):
|
||||
from tools import ReadFileTool
|
||||
f = tmp_path / "f.txt"
|
||||
f.write_text("line1\nline2\n", encoding="utf-8")
|
||||
tool = ReadFileTool()
|
||||
r = tool.execute(filepath=str(f), start_line="1'; rm -rf /; echo '",
|
||||
end_line=5)
|
||||
assert "参数错误" in r
|
||||
|
||||
def test_valid_integers_still_work(self, tmp_path):
|
||||
from tools import ReadFileTool
|
||||
f = tmp_path / "f.txt"
|
||||
f.write_text("line1\nline2\n", encoding="utf-8")
|
||||
tool = ReadFileTool()
|
||||
r = tool.execute(filepath=str(f), start_line=1, end_line=2)
|
||||
assert "sed -n" in r and "参数错误" not in r
|
||||
|
||||
|
||||
# ==================== L2:set_proxy 落盘失败如实报告 ====================
|
||||
|
||||
class TestL2ProxySaveFailure:
|
||||
def test_set_proxy_returns_false_on_save_failure(self, tmp_path):
|
||||
from common import ConfigManager
|
||||
cm = ConfigManager()
|
||||
cm._config_path = str(tmp_path / "config.json")
|
||||
# J5(审计 2026-08-19):写路径改为原子写(mkstemp+os.replace),
|
||||
# 磁盘失败改由 patch 原子写函数模拟(builtins.open 已不在路径上)
|
||||
with mock.patch("common.config.atomic_write_text",
|
||||
side_effect=PermissionError("disk ro")):
|
||||
ok = cm.set_proxy("127.0.0.1", 1080, save=True)
|
||||
assert ok is False
|
||||
|
||||
def test_proxy_tool_reports_save_failure(self, tmp_path):
|
||||
from tools import SetProxyTool
|
||||
tool = SetProxyTool(None)
|
||||
# J5:磁盘失败经原子写函数模拟
|
||||
with mock.patch("common.config.atomic_write_text",
|
||||
side_effect=PermissionError("disk ro")):
|
||||
r = tool.execute(host="127.0.0.1", port=1080, save=True,
|
||||
skip_check=True, enabled=True)
|
||||
assert "保存失败" in r
|
||||
|
||||
|
||||
# ==================== L3:SSE 非 UTF-8 行不中断流 ====================
|
||||
|
||||
class TestL3SseDecode:
|
||||
def test_non_utf8_line_does_not_break_stream(self):
|
||||
c = DeepSeekClient(api_key="k", model="m", base_url="https://x")
|
||||
resp = mock.MagicMock()
|
||||
resp.status_code = 200
|
||||
resp.iter_lines.return_value = [
|
||||
b"data: \xff\xfe",
|
||||
b'data: {"choices":[{"delta":{"content":"ok"}}]}',
|
||||
b"data: [DONE]",
|
||||
]
|
||||
with mock.patch("llm_client.requests.post", return_value=resp):
|
||||
events = list(c.stream_chat([{"role": "user", "content": "hi"}]))
|
||||
assert any(e.get("type") == "content" and e.get("text") == "ok"
|
||||
for e in events)
|
||||
|
||||
|
||||
# ==================== L4:persona 软链越界 ====================
|
||||
|
||||
class TestL4PersonaSymlink:
|
||||
def test_symlink_outside_skills_dir_rejected(self, tmp_path, monkeypatch):
|
||||
from test_conversation_compressor import patch_const
|
||||
from aibox import ExpressEmotionTool
|
||||
skills_dir = tmp_path / "skills"
|
||||
skills_dir.mkdir()
|
||||
outside = tmp_path / "outside.md"
|
||||
outside.write_text("秘密内容", encoding="utf-8")
|
||||
(skills_dir / "evil.md").symlink_to(outside)
|
||||
tool = ExpressEmotionTool()
|
||||
patch_const(monkeypatch, "SKILLS_DIR", str(skills_dir))
|
||||
content = tool._load_persona("evil", "")
|
||||
assert content == ""
|
||||
|
||||
def test_normal_persona_still_loads(self, tmp_path, monkeypatch):
|
||||
from test_conversation_compressor import patch_const
|
||||
from aibox import ExpressEmotionTool
|
||||
skills_dir = tmp_path / "skills"
|
||||
skills_dir.mkdir()
|
||||
(skills_dir / "happy.md").write_text("快乐人格", encoding="utf-8")
|
||||
tool = ExpressEmotionTool()
|
||||
patch_const(monkeypatch, "SKILLS_DIR", str(skills_dir))
|
||||
content = tool._load_persona("happy", "")
|
||||
assert content == "快乐人格"
|
||||
|
||||
|
||||
# ==================== L5:截断无半个字符 ====================
|
||||
|
||||
class TestL5NoPartialChar:
|
||||
def test_truncated_output_has_no_replacement_char(self):
|
||||
from subsystems.executors import SubprocessCommandExecutor
|
||||
ex = SubprocessCommandExecutor()
|
||||
r = ex.execute("python3 -c \"print('中' * 600000)\"", timeout=60)
|
||||
assert "截断" in r
|
||||
assert "\ufffd" not in r
|
||||
|
||||
@pytest.mark.skipif(os.name != "posix", reason="Linux 进程组语义")
|
||||
def test_grandchild_holding_pipe_is_reaped(self):
|
||||
"""`(sleep N &) ; echo done` 形态:shell 退出后孙进程持管道写端,
|
||||
宽限窗口到期强杀进程组——不悬挂 reader、孙进程被清理"""
|
||||
import subprocess as sp
|
||||
from subsystems.executors import SubprocessCommandExecutor
|
||||
marker = f"sleep {12340}"
|
||||
ex = SubprocessCommandExecutor()
|
||||
r = ex.execute(f"({marker} &) ; echo done", timeout=8)
|
||||
assert "done" in r
|
||||
leftover = sp.run(["pgrep", "-f", marker],
|
||||
capture_output=True, text=True).stdout.strip()
|
||||
assert leftover == ""
|
||||
|
||||
|
||||
# ==================== L6:compressor 除零 ====================
|
||||
|
||||
class TestL6CompressorDivZero:
|
||||
def test_empty_content_compress_no_div0(self, monkeypatch):
|
||||
comp = ConversationCompressor(llm_client=FakeLLMClient(),
|
||||
logger=None)
|
||||
from test_conversation_compressor import patch_const
|
||||
patch_const(monkeypatch, "COMPRESS_MESSAGE_THRESHOLD", 10)
|
||||
messages = [{"role": "system", "content": "SYS"}] + [
|
||||
{"role": "user" if i % 2 == 0 else "assistant", "content": ""}
|
||||
for i in range(15)
|
||||
]
|
||||
result = comp.compress_conversation(messages)
|
||||
assert isinstance(result, list)
|
||||
|
||||
|
||||
# ==================== L7:InstallSkillTool 路径引用 ====================
|
||||
|
||||
class TestL7SkillPathQuoting:
|
||||
def test_guide_commands_quote_paths(self):
|
||||
from tools import InstallSkillTool
|
||||
tool = InstallSkillTool()
|
||||
r = tool.execute(skill_path="/tmp/a b; rm -rf /")
|
||||
assert "ls -l '/tmp/a b; rm -rf /'" in r
|
||||
assert "cp '/tmp/a b; rm -rf /' " in r
|
||||
|
||||
|
||||
# ==================== L8:ask_user 缺参 ====================
|
||||
|
||||
class TestL8AskUser:
|
||||
def test_missing_question_returns_error(self):
|
||||
from tools import AskUserTool
|
||||
tool = AskUserTool()
|
||||
assert "需要提供问题" in tool.execute()
|
||||
|
||||
def test_question_passed_through(self):
|
||||
from tools import AskUserTool
|
||||
tool = AskUserTool()
|
||||
assert "测试问题" in tool.execute(question="测试问题")
|
||||
|
||||
|
||||
# ==================== L10:web 工具 scheme 白名单 ====================
|
||||
|
||||
class TestL10WebScheme:
|
||||
def test_file_url_rejected(self):
|
||||
from tools import SimpleWebCrawlerTool
|
||||
tool = SimpleWebCrawlerTool()
|
||||
r = tool.execute(action="get", url="file:///etc/shadow")
|
||||
assert "仅支持 http/https" in r
|
||||
|
||||
def test_http_url_accepted(self):
|
||||
from tools import SimpleWebCrawlerTool
|
||||
tool = SimpleWebCrawlerTool()
|
||||
r = tool.execute(action="get", url="https://example.com")
|
||||
assert "curl" in r and "仅支持" not in r
|
||||
|
||||
|
||||
# ==================== L11:registry 打码包含匹配 ====================
|
||||
|
||||
class TestL11SanitizeSubstring:
|
||||
def test_variant_keys_redacted(self):
|
||||
from tools.registry import _sanitize_kwargs
|
||||
out = _sanitize_kwargs({
|
||||
"apiKey": "x", "access_token": "y", "passwd": "z",
|
||||
"url": "https://ok",
|
||||
})
|
||||
assert out == {"apiKey": "***", "access_token": "***",
|
||||
"passwd": "***", "url": "https://ok"}
|
||||
|
||||
def test_normal_keys_untouched(self):
|
||||
from tools.registry import _sanitize_kwargs
|
||||
out = _sanitize_kwargs({"filepath": "/tmp/a", "question": "hi"})
|
||||
assert out == {"filepath": "/tmp/a", "question": "hi"}
|
||||
|
||||
|
||||
# ==================== L12:failover 日志 Bearer 脱敏 ====================
|
||||
|
||||
class TestL12LogSanitize:
|
||||
def test_sanitize_msg_redacts_bearer(self):
|
||||
msg = "瞬时故障: HTTP 401: Bearer sk-abc123xyz"
|
||||
out = _sanitize_log_msg(msg)
|
||||
assert "sk-abc123xyz" not in out
|
||||
assert "Bearer ***" in out
|
||||
|
||||
def test_failover_log_redacts_bearer(self):
|
||||
logger = mock.MagicMock()
|
||||
c = FailoverClient([], logger=logger)
|
||||
c._log("warning", "瞬时故障(Bearer sk-secret456)")
|
||||
logged = logger.warning.call_args[0][0]
|
||||
assert "sk-secret456" not in logged
|
||||
assert "Bearer ***" in logged
|
||||
@@ -0,0 +1,237 @@
|
||||
"""aibox 第四轮审计 Batch I 修复回归测试(2026-08-18,审计报告 §十二)
|
||||
|
||||
覆盖:
|
||||
- N9:prompt 写入内容长度上限(持久化注入面约束)
|
||||
- Q7b:load_conversation 越界路径拒绝(realpath 前缀校验)
|
||||
- F5:run_light_subagent 取消传播(流式可中断 + 父已取消早退)
|
||||
- Q5:提醒时间格式异常的任务按到期返回(不再静默丢弃)
|
||||
- Q8:对话文件原子写(无 .tmp 残留、JSON 完整)
|
||||
- F6:WebSocketHandler 命令专用线程池存在且可释放
|
||||
- limit=0 语义修正(None 才不加 LIMIT)
|
||||
|
||||
Q4(busy_timeout 10s)为时序性行为(默认 5s 等待已能通过短竞争),
|
||||
不做慢速单测;F9(--once 硬超时)为进程级行为,由 CLI 集成面覆盖。
|
||||
"""
|
||||
|
||||
import os
|
||||
import queue
|
||||
import sqlite3
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import types
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
# ---- 环境桩:websockets 未安装时保证 import aibox.py 不失败 ----
|
||||
if "websockets" not in sys.modules:
|
||||
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||
|
||||
import aibox # noqa: E402
|
||||
from aibox import run_light_subagent # noqa: E402
|
||||
|
||||
from subsystems.db import MemoryDatabase, MemoDatabase # noqa: E402
|
||||
from subsystems.prompt_manager import PromptManager # noqa: E402
|
||||
from subsystems.conversation_history import ConversationHistory # noqa: E402
|
||||
from tools.prompt_tools import ( # noqa: E402
|
||||
UpdateIdentityTool,
|
||||
UpdateSoulTool,
|
||||
SaveMemoryAndEndConversationTool,
|
||||
)
|
||||
|
||||
|
||||
# ==================== N9:prompt 写入长度上限 ====================
|
||||
|
||||
class TestN9PromptContentCap:
|
||||
def test_update_identity_oversized_rejected(self, tmp_path, monkeypatch):
|
||||
"""manager 层兜底:超长身份内容拒绝写入"""
|
||||
pm = PromptManager()
|
||||
with mock.patch.object(pm, "_prompts_dir", str(tmp_path)):
|
||||
assert pm.update_identity("x" * (PromptManager.MAX_CONTENT_LENGTH + 1)) is False
|
||||
assert not (tmp_path / pm.IDENTITY_FILE).exists()
|
||||
|
||||
def test_update_identity_tool_friendly_error(self, tmp_path):
|
||||
"""工具层:超长内容返回明确提示而非写入"""
|
||||
pm = PromptManager()
|
||||
with mock.patch.object(pm, "_prompts_dir", str(tmp_path)):
|
||||
tool = UpdateIdentityTool(pm, None)
|
||||
r = tool.execute(content="x" * (PromptManager.MAX_CONTENT_LENGTH + 1))
|
||||
assert "内容过长" in r
|
||||
|
||||
def test_update_soul_tool_friendly_error(self, tmp_path):
|
||||
pm = PromptManager()
|
||||
with mock.patch.object(pm, "_prompts_dir", str(tmp_path)):
|
||||
tool = UpdateSoulTool(pm, None)
|
||||
r = tool.execute(content="y" * (PromptManager.MAX_CONTENT_LENGTH + 1))
|
||||
assert "内容过长" in r
|
||||
|
||||
def test_save_memory_status_oversized_rejected(self, tmp_path):
|
||||
"""save_memory_and_end_conversation:状态内容超长拒绝"""
|
||||
pm = PromptManager()
|
||||
with mock.patch.object(pm, "_prompts_dir", str(tmp_path)):
|
||||
tool = SaveMemoryAndEndConversationTool(
|
||||
exit_callback=lambda *a: None,
|
||||
memory_db=None, prompt_manager=pm, logger=None)
|
||||
r = tool.execute(summary="这是一段足够长的对话总结" * 20,
|
||||
current_status="x" * PromptManager.MAX_CONTENT_LENGTH)
|
||||
assert "状态内容过长" in r
|
||||
|
||||
def test_normal_content_still_written(self, tmp_path):
|
||||
pm = PromptManager()
|
||||
with mock.patch.object(pm, "_prompts_dir", str(tmp_path)):
|
||||
assert pm.update_identity("正常身份内容") is True
|
||||
assert "正常身份内容" in (tmp_path / pm.IDENTITY_FILE).read_text(encoding="utf-8")
|
||||
|
||||
|
||||
# ==================== Q7b:load_conversation 路径校验 ====================
|
||||
|
||||
class TestQ7bLoadConversationBoundary:
|
||||
def test_traversal_rejected(self, tmp_path):
|
||||
h = ConversationHistory(save_dir=str(tmp_path / "save"), load_history=False)
|
||||
secret = tmp_path / "secret.json"
|
||||
secret.write_text('{"leak": true}', encoding="utf-8")
|
||||
assert h.load_conversation("../secret.json") is None
|
||||
|
||||
def test_normal_file_still_loads(self, tmp_path):
|
||||
h = ConversationHistory(save_dir=str(tmp_path / "save"), load_history=False)
|
||||
h.current_conversation = {
|
||||
"id": "1", "start_time": "2026-08-18 00:00:00", "end_time": None,
|
||||
"topic": "t", "messages": [{"role": "user", "content": "hi"}],
|
||||
"tool_calls": [], "stats": {},
|
||||
}
|
||||
h.save_current_conversation()
|
||||
h.archive_current_conversation()
|
||||
archives = list(h.archive_dir.glob("conversation_*.json"))
|
||||
assert archives
|
||||
loaded = h.load_conversation(archives[0].name)
|
||||
assert loaded is not None and loaded["id"] == "1"
|
||||
|
||||
|
||||
# ==================== F5:run_light_subagent 取消传播 ====================
|
||||
|
||||
class _BlockingLightClient:
|
||||
"""流式产出首段后卡住,直到 cancel() 才中断的伪 LLM 客户端"""
|
||||
|
||||
model = "fake-model"
|
||||
default_temperature = None
|
||||
|
||||
def __init__(self):
|
||||
self.cancelled = threading.Event()
|
||||
|
||||
def cancel(self):
|
||||
self.cancelled.set()
|
||||
|
||||
def stream_chat(self, messages, **kwargs):
|
||||
yield {"type": "content", "text": "开头"}
|
||||
while not self.cancelled.wait(timeout=0.1):
|
||||
pass
|
||||
raise RuntimeError("连接已中断(cancel)")
|
||||
|
||||
|
||||
class TestF5LightSubagentCancel:
|
||||
def test_cancelled_before_start_returns_early(self):
|
||||
"""父事件已置位 → 直接返回取消提示,不创建客户端"""
|
||||
parent = threading.Event()
|
||||
parent.set()
|
||||
with mock.patch("aibox.create_llm_client") as mk:
|
||||
r = run_light_subagent("r", "t", None, cancel_event=parent)
|
||||
assert "取消" in r
|
||||
mk.assert_not_called()
|
||||
|
||||
def test_cancel_during_stream_aborts(self):
|
||||
"""流式中父级取消 → 观察者线程 cancel 客户端 → 立即中断返回"""
|
||||
parent = threading.Event()
|
||||
fake = _BlockingLightClient()
|
||||
with mock.patch("aibox.create_llm_client", return_value=fake):
|
||||
result_holder = {}
|
||||
|
||||
def _run():
|
||||
result_holder["r"] = run_light_subagent(
|
||||
"r", "t", None, cancel_event=parent)
|
||||
|
||||
th = threading.Thread(target=_run, daemon=True)
|
||||
th.start()
|
||||
time.sleep(0.5) # 等待首段产出、观察者就绪
|
||||
parent.set()
|
||||
th.join(timeout=5)
|
||||
assert not th.is_alive() # 未挂死
|
||||
assert "失败" in result_holder.get("r", "")
|
||||
|
||||
def test_express_emotion_registered_with_parent_event(self):
|
||||
"""DeepSeekChat 注册 ExpressEmotionTool 时注入自身取消事件"""
|
||||
from test_once_mode import _make_ai
|
||||
ai = _make_ai()
|
||||
tool = next((t for t in ai.tool_registry.get_all_tools()
|
||||
if t.get_name() == "express_emotion"), None)
|
||||
assert tool is not None
|
||||
assert tool.cancel_event is ai._cancel_event
|
||||
|
||||
|
||||
# ==================== Q5:格式异常任务不丢弃 ====================
|
||||
|
||||
class TestQ5MalformedReminderTime:
|
||||
def test_malformed_time_task_returned_as_due(self, tmp_path):
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
# J1(审计 2026-08-19):add_memo 已拒绝畸形时间(入库校验),
|
||||
# 模拟校验落地前的存量脏数据需绕过 add_memo 直接 SQL 插入
|
||||
import sqlite3
|
||||
with sqlite3.connect(str(tmp_path / "m.db")) as conn:
|
||||
cur = conn.cursor()
|
||||
cur.execute('''INSERT INTO memos
|
||||
(title, content, created_at, reminder_time, is_completed, source,
|
||||
repeat_type, repeat_interval, repeat_interval_unit, is_immediate)
|
||||
VALUES (?, ?, ?, ?, 0, 'test', 'none', 0, 'days', 0)''',
|
||||
("格式异常", "c", "2026-08-18 00:00:00", "不是时间"))
|
||||
mid = cur.lastrowid
|
||||
pending = db.get_pending_memos(before_time="2026-08-18 00:00:00")
|
||||
assert any(m["id"] == mid for m in pending)
|
||||
|
||||
def test_normal_pending_still_filtered(self, tmp_path):
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
db.add_memo(title="未来任务", content="c", reminder_time="2099-01-01 00:00:00")
|
||||
pending = db.get_pending_memos(before_time="2026-08-18 00:00:00")
|
||||
assert pending == []
|
||||
|
||||
|
||||
# ==================== Q8:对话文件原子写 ====================
|
||||
|
||||
class TestQ8AtomicWrites:
|
||||
def test_no_tmp_leftover_and_valid_json(self, tmp_path):
|
||||
h = ConversationHistory(save_dir=str(tmp_path / "save"), load_history=False)
|
||||
h.current_conversation = {
|
||||
"id": "1", "start_time": "2026-08-18 00:00:00", "end_time": None,
|
||||
"topic": "t", "messages": [], "tool_calls": [], "stats": {},
|
||||
}
|
||||
h.save_current_conversation()
|
||||
assert h.current_file.exists() # 文件名可能被用户配置覆盖,用实例属性
|
||||
import json
|
||||
json.loads(h.current_file.read_text(encoding="utf-8")) # 不抛即合法 JSON
|
||||
h.archive_current_conversation() # 归档后 current 文件按设计移除
|
||||
assert list(Path(tmp_path).rglob("*.tmp")) == []
|
||||
|
||||
|
||||
# ==================== F6:WS 命令专用线程池 ====================
|
||||
|
||||
class TestF6CmdExecutor:
|
||||
def test_handler_has_dedicated_executor(self):
|
||||
from aibox import WebSocketHandler
|
||||
h = WebSocketHandler(ai_manager=None)
|
||||
assert h._cmd_executor is not None
|
||||
h.shutdown()
|
||||
assert h._cmd_executor is None
|
||||
|
||||
|
||||
# ==================== limit=0 语义 ====================
|
||||
|
||||
class TestLimitZeroSemantics:
|
||||
def test_pending_memos_limit_zero_returns_none(self, tmp_path):
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
db.add_memo(title="t", content="c", reminder_time="2026-01-01 00:00:00")
|
||||
assert db.get_pending_memos(limit=0, before_time="2026-08-18 00:00:00") == []
|
||||
|
||||
def test_immediate_memos_limit_zero_returns_none(self, tmp_path):
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
db.add_memo(title="t", content="c", is_immediate=True)
|
||||
assert db.get_immediate_memos(limit=0) == []
|
||||
@@ -0,0 +1,540 @@
|
||||
"""aibox 第五轮审计(Batch K)修复回归测试(2026-08-19)
|
||||
|
||||
波次 1:
|
||||
- J3:FTS 外部内容表删除换序后索引无残留(delete_memory / delete_memories_by_ids)
|
||||
- J4:model_config 连通性验证异常消息 Bearer 令牌脱敏
|
||||
- J6:PromptManager/ConfigManager 单例并发初始化不双重加载
|
||||
- J8:add_profile 对非法温度/超时的容错(交互层不再抢先 float()/int())
|
||||
|
||||
波次 2(J1 任务系统投毒):
|
||||
- add_memo 入库校验:畸形时间/重复间隔/重复类型拒绝,三入口一处覆盖
|
||||
- 存量脏数据防御:调度器预检、mark_as_triggered、展示路径不再打崩
|
||||
- WS task add 友好错误 + register 握手异常可回收
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import threading
|
||||
import types
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
if "websockets" not in sys.modules:
|
||||
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||
|
||||
import aibox # noqa: E402
|
||||
import model_config # noqa: E402
|
||||
from common import ConfigManager # noqa: E402
|
||||
from subsystems.db import MemoDatabase, MemoryDatabase # noqa: E402
|
||||
from subsystems.prompt_manager import PromptManager # noqa: E402
|
||||
|
||||
from test_model_config import _cfg # noqa: E402
|
||||
|
||||
|
||||
# ==================== J3:FTS 删除顺序 ====================
|
||||
|
||||
def _fts_hits(db_path, keyword):
|
||||
conn = sqlite3.connect(str(db_path))
|
||||
try:
|
||||
return conn.execute(
|
||||
"SELECT rowid FROM memories_fts WHERE memories_fts MATCH ?",
|
||||
(f'"{keyword}"',)).fetchall()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _add_mem(db, content):
|
||||
return db.add_memory(topic="话题", summary="摘要", content=content,
|
||||
start_time="2026-08-19 00:00:00",
|
||||
end_time="2026-08-19 00:00:01", importance=2)
|
||||
|
||||
|
||||
class TestJ3FtsDeleteOrder:
|
||||
def test_delete_memory_leaves_no_fts_orphan(self, tmp_path):
|
||||
"""单条删除后 FTS 索引不得残留(修复前 MATCH 仍命中已删 rowid)"""
|
||||
db = MemoryDatabase(db_path=str(tmp_path / "mem.db"))
|
||||
mid = _add_mem(db, "机密内容甲")
|
||||
assert _fts_hits(tmp_path / "mem.db", "机密内容甲")
|
||||
assert db.delete_memory(mid) is True
|
||||
assert _fts_hits(tmp_path / "mem.db", "机密内容甲") == []
|
||||
|
||||
def test_batch_delete_leaves_no_fts_orphan(self, tmp_path):
|
||||
"""批量删除同样不留索引残留"""
|
||||
db = MemoryDatabase(db_path=str(tmp_path / "mem.db"))
|
||||
ids = [_add_mem(db, f"机密内容乙{i}") for i in range(3)]
|
||||
assert db.delete_memories_by_ids(ids) == 3
|
||||
for i in range(3):
|
||||
assert _fts_hits(tmp_path / "mem.db", f"机密内容乙{i}") == []
|
||||
|
||||
def test_search_after_delete_no_hit(self, tmp_path):
|
||||
"""删除后搜索(FTS 路径)不再命中已删内容"""
|
||||
db = MemoryDatabase(db_path=str(tmp_path / "mem.db"))
|
||||
mid = _add_mem(db, "独特关键词丙")
|
||||
assert db.search_memories(["独特关键词丙"])
|
||||
db.delete_memory(mid)
|
||||
assert db.search_memories(["独特关键词丙"]) == []
|
||||
|
||||
|
||||
# ==================== J4:验证异常脱敏 ====================
|
||||
|
||||
class TestJ4TestProfileSanitize:
|
||||
def _models_cfg(self, tmp_path):
|
||||
return _cfg(tmp_path, {"models": {"bad": {
|
||||
"api_url": "https://api.example.com/v1/chat/completions",
|
||||
"model": "m", "api_key_env": "K"}}})
|
||||
|
||||
def test_bearer_token_masked_in_failure_message(self, tmp_path):
|
||||
"""网关错误页回显 Authorization 头时,令牌不得出现在终端回显"""
|
||||
cm = self._models_cfg(tmp_path)
|
||||
err = RuntimeError("HTTP 502: Authorization: Bearer sk-secret12345")
|
||||
with mock.patch("llm_client.create_llm_client", side_effect=err):
|
||||
result = model_config.test_profile(cm, "bad")
|
||||
assert "sk-secret12345" not in result
|
||||
assert "Bearer ***" in result
|
||||
|
||||
def test_plain_error_message_preserved(self, tmp_path):
|
||||
"""无敏感内容的错误消息原样保留(脱敏不误伤)"""
|
||||
cm = self._models_cfg(tmp_path)
|
||||
with mock.patch("llm_client.create_llm_client",
|
||||
side_effect=RuntimeError("连接被拒绝")):
|
||||
result = model_config.test_profile(cm, "bad")
|
||||
assert "连接被拒绝" in result
|
||||
|
||||
|
||||
# ==================== J6:单例并发初始化 ====================
|
||||
|
||||
class TestJ6SingletonInitRace:
|
||||
def test_prompt_manager_concurrent_init_no_duplication(self, tmp_path):
|
||||
"""多线程首例并发:单实例且提示词只加载一份(修复前整份翻倍)"""
|
||||
PromptManager._instance = None
|
||||
PromptManager._prompts = []
|
||||
PromptManager._prompts_dir = None
|
||||
prompts_dir = tmp_path / "prompts"
|
||||
prompts_dir.mkdir()
|
||||
try:
|
||||
barrier = threading.Barrier(8)
|
||||
instances = [None] * 8
|
||||
|
||||
def worker(k):
|
||||
barrier.wait()
|
||||
instances[k] = PromptManager()
|
||||
|
||||
with mock.patch.object(ConfigManager, "get",
|
||||
return_value=str(prompts_dir)):
|
||||
threads = [threading.Thread(target=worker, args=(k,))
|
||||
for k in range(8)]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
assert len({id(i) for i in instances}) == 1
|
||||
files = [p["file"] for p in PromptManager._prompts]
|
||||
assert files, "提示词应已加载"
|
||||
assert len(files) == len(set(files)), f"提示词被重复加载: {files}"
|
||||
finally:
|
||||
PromptManager._instance = None
|
||||
PromptManager._prompts = []
|
||||
PromptManager._prompts_dir = None
|
||||
|
||||
def test_config_manager_concurrent_init_single_load(self, tmp_path):
|
||||
"""多线程首例并发:单实例且 _load_config 只执行一次"""
|
||||
cfg_file = tmp_path / "config.json"
|
||||
calls = {"n": 0}
|
||||
original_load = ConfigManager._load_config
|
||||
|
||||
def counting_load(self):
|
||||
calls["n"] += 1
|
||||
original_load(self)
|
||||
|
||||
ConfigManager._instance = None
|
||||
ConfigManager._config = None
|
||||
ConfigManager._config_path = None
|
||||
try:
|
||||
barrier = threading.Barrier(8)
|
||||
instances = [None] * 8
|
||||
|
||||
def worker(k):
|
||||
barrier.wait()
|
||||
instances[k] = ConfigManager()
|
||||
|
||||
with mock.patch.dict(os.environ, {"AI_CONFIG_PATH": str(cfg_file)}), \
|
||||
mock.patch.object(ConfigManager, "_load_config", counting_load):
|
||||
threads = [threading.Thread(target=worker, args=(k,))
|
||||
for k in range(8)]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
assert len({id(i) for i in instances}) == 1
|
||||
assert calls["n"] == 1
|
||||
finally:
|
||||
ConfigManager._instance = None
|
||||
ConfigManager._config = None
|
||||
ConfigManager._config_path = None
|
||||
|
||||
|
||||
# ==================== J8:交互输入容错 ====================
|
||||
|
||||
class TestJ8NumericTolerance:
|
||||
def test_invalid_temperature_returns_friendly_error(self, tmp_path):
|
||||
"""非数字温度 → 友好错误而非 ValueError(交互层不再抢先 float())"""
|
||||
cm = _cfg(tmp_path, {"models": {}})
|
||||
msgs = model_config.add_profile(
|
||||
cm, "p1", "https://api.example.com/v1/chat/completions", "m1",
|
||||
api_key_env="K", temperature="not-a-number")
|
||||
assert any(m.startswith("❌") and "温度" in m for m in msgs)
|
||||
assert "p1" not in model_config.get_models(cm)
|
||||
|
||||
def test_invalid_timeout_returns_friendly_error(self, tmp_path):
|
||||
cm = _cfg(tmp_path, {"models": {}})
|
||||
msgs = model_config.add_profile(
|
||||
cm, "p1", "https://api.example.com/v1/chat/completions", "m1",
|
||||
api_key_env="K", timeout="abc")
|
||||
assert any(m.startswith("❌") and "超时" in m for m in msgs)
|
||||
assert "p1" not in model_config.get_models(cm)
|
||||
|
||||
def test_string_numbers_accepted_and_clamped(self, tmp_path):
|
||||
"""数字字符串正常接受并钳制(交互层原始字符串直传语义)"""
|
||||
cm = _cfg(tmp_path, {"models": {}})
|
||||
msgs = model_config.add_profile(
|
||||
cm, "p1", "https://api.example.com/v1/chat/completions", "m1",
|
||||
api_key_env="K", temperature="0.7", timeout="120")
|
||||
assert all(not m.startswith("❌") for m in msgs)
|
||||
profile = model_config.get_models(cm)["p1"]
|
||||
assert profile["temperature"] == 0.7
|
||||
assert profile["timeout"] == 120
|
||||
|
||||
|
||||
# ==================== J1 波次 2:任务系统投毒 ====================
|
||||
|
||||
def _insert_memo_raw(db_path, reminder_time="2026-08-18 00:00:00",
|
||||
repeat_type="none", repeat_interval=0,
|
||||
repeat_interval_unit="days", repeat_end_time=None):
|
||||
"""直接 SQL 插入任务行——模拟校验落地前的存量脏数据(绕过 add_memo)"""
|
||||
with sqlite3.connect(str(db_path)) as conn:
|
||||
cur = conn.cursor()
|
||||
cur.execute('''INSERT INTO memos
|
||||
(title, content, created_at, reminder_time, tags, priority, created_by,
|
||||
is_completed, source, repeat_type, repeat_interval, repeat_interval_unit,
|
||||
repeat_end_time, is_immediate)
|
||||
VALUES ('t', 'c', '2026-08-18 00:00:00', ?, '', 1, 'test',
|
||||
0, 'test', ?, ?, ?, ?, 0)''',
|
||||
(reminder_time, repeat_type, repeat_interval, repeat_interval_unit,
|
||||
repeat_end_time))
|
||||
return cur.lastrowid
|
||||
|
||||
|
||||
class TestJ1AddMemoValidation:
|
||||
def test_malformed_reminder_time_rejected(self, tmp_path):
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
with pytest.raises(ValueError, match="提醒时间格式非法"):
|
||||
db.add_memo(title="t", reminder_time="不是时间")
|
||||
with sqlite3.connect(str(tmp_path / "m.db")) as conn:
|
||||
assert conn.execute("SELECT COUNT(*) FROM memos").fetchone()[0] == 0
|
||||
|
||||
def test_malformed_repeat_end_time_rejected(self, tmp_path):
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
with pytest.raises(ValueError, match="重复结束时间格式非法"):
|
||||
db.add_memo(title="t", reminder_time="2026-08-18 00:00:00",
|
||||
repeat_type="daily", repeat_end_time="不是时间")
|
||||
|
||||
def test_invalid_repeat_type_rejected(self, tmp_path):
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
with pytest.raises(ValueError, match="无效的重复类型"):
|
||||
db.add_memo(title="t", reminder_time="2026-08-18 00:00:00",
|
||||
repeat_type="every-second")
|
||||
|
||||
def test_non_numeric_repeat_interval_rejected(self, tmp_path):
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
with pytest.raises(ValueError, match="重复间隔必须是整数"):
|
||||
db.add_memo(title="t", reminder_time="2026-08-18 00:00:00",
|
||||
repeat_type="custom", repeat_interval="垃圾")
|
||||
|
||||
def test_custom_repeat_requires_positive_interval(self, tmp_path):
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
with pytest.raises(ValueError, match="repeat_interval"):
|
||||
db.add_memo(title="t", reminder_time="2026-08-18 00:00:00",
|
||||
repeat_type="custom", repeat_interval=0)
|
||||
|
||||
def test_numeric_string_interval_coerced(self, tmp_path):
|
||||
"""数字字符串被强转接受(WS 传 JSON 数字/字符串均安全)"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
mid = db.add_memo(title="t", reminder_time="2026-08-18 00:00:00",
|
||||
repeat_type="custom", repeat_interval="30",
|
||||
repeat_interval_unit="minutes")
|
||||
assert db.get_memo(mid)["repeat_interval"] == 30
|
||||
|
||||
def test_add_user_task_propagates_validation(self, tmp_path):
|
||||
"""WS 路径(add_user_task 薄包装 add_memo)同样被校验拦截"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
with pytest.raises(ValueError, match="提醒时间格式非法"):
|
||||
db.add_user_task(title="t", reminder_time="不是时间")
|
||||
|
||||
|
||||
class TestJ1DirtyDataDefense:
|
||||
def test_mark_as_triggered_completes_malformed_one_shot(self, tmp_path):
|
||||
"""存量畸形 reminder_time 一次性任务:触发不再抛异常,正常完成"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
mid = _insert_memo_raw(tmp_path / "m.db", reminder_time="不是时间")
|
||||
assert db.mark_as_triggered(mid) is None
|
||||
assert db.get_memo(mid)["is_completed"] == 1
|
||||
|
||||
def test_mark_as_triggered_handles_garbage_repeat_interval(self, tmp_path):
|
||||
"""存量 custom 任务 repeat_interval 为文本:按非重复处理,不打崩事务"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
mid = _insert_memo_raw(tmp_path / "m.db", reminder_time="2026-08-18 00:00:00",
|
||||
repeat_type="custom", repeat_interval="垃圾")
|
||||
assert db.mark_as_triggered(mid) is None
|
||||
assert db.get_memo(mid)["is_completed"] == 1
|
||||
|
||||
def test_malformed_repeat_end_time_treated_as_no_end(self, tmp_path):
|
||||
"""存量畸形 repeat_end_time:解析失败视为无结束时间,不抛异常"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
mid = _insert_memo_raw(tmp_path / "m.db", reminder_time="2026-08-19 00:00:00",
|
||||
repeat_type="daily", repeat_end_time="不是时间")
|
||||
memo = db.get_memo(mid)
|
||||
from datetime import datetime
|
||||
base = datetime.strptime("2026-08-19 00:00:00", "%Y-%m-%d %H:%M:%S")
|
||||
assert db._should_end_repeat(memo, base) is False
|
||||
assert db._calculate_next_reminder(memo, base) is not None
|
||||
|
||||
def test_pending_tasks_for_user_malformed_time_overdue(self, tmp_path):
|
||||
"""展示路径:畸形 reminder_time 行按到期展示,不再打崩握手"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
mid = _insert_memo_raw(tmp_path / "m.db", reminder_time="不是时间")
|
||||
tasks = db.get_pending_tasks_for_user()
|
||||
row = next(t for t in tasks if t["id"] == mid)
|
||||
assert row["is_overdue"] is True
|
||||
|
||||
|
||||
class TestJ1WsAndScheduler:
|
||||
def test_ws_task_add_rejects_malformed_repeat_end(self, tmp_path):
|
||||
"""WS task add 传畸形 repeat_end_time:友好错误且不入库"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
ai_manager = mock.MagicMock()
|
||||
ai_manager.memo_db = db
|
||||
handler = aibox.WebSocketHandler(ai_manager)
|
||||
try:
|
||||
ws = mock.AsyncMock()
|
||||
asyncio.run(handler.handle_task_message(ws, {
|
||||
"type": "task", "action": "add", "title": "t", "delay": 5,
|
||||
"repeat_type": "daily", "repeat_end_time": "不是时间"}))
|
||||
sent = json.loads(ws.send.call_args[0][0])
|
||||
assert sent["type"] == "error"
|
||||
assert "任务添加失败" in sent["content"]
|
||||
with sqlite3.connect(str(tmp_path / "m.db")) as conn:
|
||||
assert conn.execute("SELECT COUNT(*) FROM memos").fetchone()[0] == 0
|
||||
finally:
|
||||
handler.shutdown()
|
||||
|
||||
def test_scheduler_precheck_malformed_time_executes_not_loops(self, tmp_path):
|
||||
"""调度器预检:畸形 reminder_time 按到期执行并标记触发(不再热循环)"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
mid = _insert_memo_raw(tmp_path / "m.db", reminder_time="不是时间")
|
||||
loop = asyncio.new_event_loop()
|
||||
executed = {"n": 0}
|
||||
try:
|
||||
scheduler = aibox.TaskScheduler(loop=loop, logger=None)
|
||||
scheduler.broadcast_enabled = False
|
||||
|
||||
async def fake_execute(memo, memo_db):
|
||||
executed["n"] += 1
|
||||
return "ok"
|
||||
|
||||
scheduler.execute_task_as_conversation = fake_execute
|
||||
loop.run_until_complete(scheduler.check_and_execute_tasks(db))
|
||||
finally:
|
||||
loop.close()
|
||||
assert executed["n"] == 1
|
||||
assert db.get_memo(mid)["is_completed"] == 1
|
||||
|
||||
|
||||
# ==================== J5 波次 4:原子写 + 损坏备份 ====================
|
||||
|
||||
class TestJ5AtomicWrites:
|
||||
def test_save_config_atomic_no_tmp_leftover(self, tmp_path):
|
||||
"""_save_config 原子写:无 .aibox_tmp_ 残留且 JSON 合法"""
|
||||
cm = ConfigManager()
|
||||
cm._config_path = str(tmp_path / "config.json")
|
||||
cm._config = {"api": {"deepseek_model": "m"}, "models": {}}
|
||||
assert cm._save_config() is True
|
||||
import json as _json
|
||||
data = _json.loads((tmp_path / "config.json").read_text(encoding="utf-8"))
|
||||
assert data["api"]["deepseek_model"] == "m"
|
||||
assert list(tmp_path.glob(".aibox_tmp_*")) == []
|
||||
|
||||
def test_corrupted_config_backed_up_before_fallback(self, tmp_path):
|
||||
"""损坏配置加载:备份 .corrupt.* 文件后回退默认,原内容不丢"""
|
||||
cfg_file = tmp_path / "config.json"
|
||||
cfg_file.write_text('{"api": {"broken": ', encoding="utf-8") # 半截 JSON
|
||||
ConfigManager._instance = None
|
||||
ConfigManager._config = None
|
||||
ConfigManager._config_path = None
|
||||
try:
|
||||
with mock.patch.dict(os.environ, {"AI_CONFIG_PATH": str(cfg_file)}):
|
||||
cm = ConfigManager()
|
||||
assert cm._config is not None
|
||||
backups = list(tmp_path.glob("config.json.corrupt.*"))
|
||||
assert len(backups) == 1
|
||||
assert backups[0].read_text(encoding="utf-8") == '{"api": {"broken": '
|
||||
finally:
|
||||
ConfigManager._instance = None
|
||||
ConfigManager._config = None
|
||||
ConfigManager._config_path = None
|
||||
|
||||
def test_prompt_write_atomic(self, tmp_path):
|
||||
"""prompt 更新原子写:内容生效且无临时文件残留"""
|
||||
from subsystems.prompt_manager import PromptManager
|
||||
PromptManager._instance = None
|
||||
PromptManager._prompts = []
|
||||
PromptManager._prompts_dir = None
|
||||
prompts_dir = tmp_path / "prompts"
|
||||
prompts_dir.mkdir()
|
||||
try:
|
||||
with mock.patch.object(ConfigManager, "get",
|
||||
return_value=str(prompts_dir)):
|
||||
pm = PromptManager()
|
||||
assert pm.update_identity("新身份内容") is True
|
||||
identity = prompts_dir / "00_IDENTITY.md"
|
||||
assert "新身份内容" in identity.read_text(encoding="utf-8")
|
||||
assert list(tmp_path.glob(".aibox_tmp_*")) == []
|
||||
finally:
|
||||
PromptManager._instance = None
|
||||
PromptManager._prompts = []
|
||||
PromptManager._prompts_dir = None
|
||||
|
||||
|
||||
# ==================== J7 波次 4:broadcast 超时回收 ====================
|
||||
|
||||
class TestJ7BroadcastTimeout:
|
||||
def test_stuck_client_timed_out_and_unregistered(self):
|
||||
"""半开 TCP 僵死客户端:broadcast 在超时内返回并回收该连接"""
|
||||
from aibox import WebSocketHandler
|
||||
handler = WebSocketHandler(ai_manager=mock.MagicMock())
|
||||
handler.broadcast_send_timeout = 0.1
|
||||
try:
|
||||
ws = mock.AsyncMock()
|
||||
|
||||
async def stuck_send(*args, **kwargs):
|
||||
await asyncio.sleep(30)
|
||||
|
||||
ws.send.side_effect = stuck_send
|
||||
handler.connected_clients.add(ws)
|
||||
asyncio.run(handler.broadcast({"type": "ping"}))
|
||||
assert ws not in handler.connected_clients
|
||||
finally:
|
||||
handler.shutdown()
|
||||
|
||||
def test_healthy_client_stays_connected(self):
|
||||
"""正常客户端:发送成功且不被回收"""
|
||||
from aibox import WebSocketHandler
|
||||
handler = WebSocketHandler(ai_manager=mock.MagicMock())
|
||||
handler.broadcast_send_timeout = 1.0
|
||||
try:
|
||||
ws = mock.AsyncMock()
|
||||
handler.connected_clients.add(ws)
|
||||
asyncio.run(handler.broadcast({"type": "ping"}))
|
||||
assert ws in handler.connected_clients
|
||||
assert ws.send.called
|
||||
finally:
|
||||
handler.shutdown()
|
||||
|
||||
|
||||
# ==================== J9 波次 5:MemoDatabase 连接层 ====================
|
||||
|
||||
class TestJ9MemoDbConnect:
|
||||
def test_wal_enabled_after_init(self, tmp_path):
|
||||
"""WAL 显式设置——不再依赖与 MemoryDatabase 的创建顺序"""
|
||||
MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
with sqlite3.connect(str(tmp_path / "m.db")) as conn:
|
||||
mode = conn.execute("PRAGMA journal_mode").fetchone()[0]
|
||||
assert mode == "wal"
|
||||
|
||||
def test_connection_factory_sets_busy_timeout(self, tmp_path):
|
||||
"""统一连接工厂:busy_timeout 覆盖全部写路径"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
conn = db._connect()
|
||||
try:
|
||||
assert conn.execute("PRAGMA busy_timeout").fetchone()[0] == 10000
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
# ==================== J10 波次 5:memo 批量删除 ====================
|
||||
|
||||
class TestJ10BatchDelete:
|
||||
def test_batch_delete_via_memo_tool(self, tmp_path):
|
||||
"""schema 声明的 memo_ids 批量删除真正生效(此前静默忽略)"""
|
||||
from tools.db_tools import MemoTool
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
ids = [db.add_memo(title=f"任务{i}",
|
||||
reminder_time="2026-08-18 00:00:00") for i in range(3)]
|
||||
tool = MemoTool(db)
|
||||
r = tool.execute(action="delete", confirm=True, memo_ids=ids)
|
||||
assert "3/3" in r
|
||||
with sqlite3.connect(str(tmp_path / "m.db")) as conn:
|
||||
assert conn.execute("SELECT COUNT(*) FROM memos").fetchone()[0] == 0
|
||||
|
||||
def test_batch_delete_requires_list(self, tmp_path):
|
||||
from tools.db_tools import MemoTool
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
tool = MemoTool(db)
|
||||
r = tool.execute(action="delete", confirm=True, memo_ids="1")
|
||||
assert "列表" in r
|
||||
|
||||
|
||||
# ==================== J15 波次 5:未知 task action ====================
|
||||
|
||||
class TestJ15UnknownTaskAction:
|
||||
def test_unknown_action_returns_error(self):
|
||||
"""未知 task action 明确报错(此前静默吞掉)"""
|
||||
from aibox import WebSocketHandler
|
||||
handler = WebSocketHandler(ai_manager=mock.MagicMock())
|
||||
try:
|
||||
ws = mock.AsyncMock()
|
||||
asyncio.run(handler.handle_task_message(ws, {"action": "explode"}))
|
||||
sent = json.loads(ws.send.call_args[0][0])
|
||||
assert sent["type"] == "error"
|
||||
assert "未知的任务操作" in sent["content"]
|
||||
finally:
|
||||
handler.shutdown()
|
||||
|
||||
|
||||
# ==================== J16 波次 5:skill_manager ====================
|
||||
|
||||
class TestJ16SkillManager:
|
||||
def test_markdown_skipped_when_executable_same_name(self, tmp_path, monkeypatch):
|
||||
"""同名键冲突:可执行技能确定性优先(此前由 listdir 顺序决定)"""
|
||||
from common import constants
|
||||
from subsystems.skill_manager import SkillManager
|
||||
skills_dir = tmp_path / "skills"
|
||||
skills_dir.mkdir()
|
||||
exe = skills_dir / "foo"
|
||||
exe.write_text("#!/bin/sh\necho '可执行技能简介'", encoding="utf-8")
|
||||
exe.chmod(0o755)
|
||||
(skills_dir / "foo.md").write_text("# Markdown 简介\n内容", encoding="utf-8")
|
||||
monkeypatch.setattr(type(constants), "SKILLS_DIR",
|
||||
property(lambda self: str(skills_dir)))
|
||||
sm = SkillManager()
|
||||
assert "foo" in sm.skills
|
||||
assert sm.skills["foo"]["is_executable"] is True
|
||||
|
||||
def test_help_output_truncated(self, tmp_path, monkeypatch):
|
||||
"""--help 输出截断(内容进系统提示词,防撑爆上下文)"""
|
||||
from common import constants
|
||||
from subsystems.skill_manager import SkillManager
|
||||
skills_dir = tmp_path / "skills"
|
||||
skills_dir.mkdir()
|
||||
exe = skills_dir / "big"
|
||||
exe.write_text("#!/bin/sh\necho '" + "x" * 5000 + "'", encoding="utf-8")
|
||||
exe.chmod(0o755)
|
||||
monkeypatch.setattr(type(constants), "SKILLS_DIR",
|
||||
property(lambda self: str(skills_dir)))
|
||||
sm = SkillManager()
|
||||
assert len(sm.skills["big"]["content"]) <= SkillManager.MAX_HELP_LENGTH + 20
|
||||
@@ -0,0 +1,157 @@
|
||||
"""aibox 第六轮审计(Batch L)修复回归测试(2026-08-20)
|
||||
|
||||
R1:/memos WS 命令对畸形 reminder_time 不再整体报错(J1 加固漏点)
|
||||
R2:handle_interrupt 只取消本会话 AI,不再触碰 CLI 主 AI 的全局 executor
|
||||
R3:unregister 补齐连接回收(关闭 websocket + 取消在途处理任务)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import sys
|
||||
import types
|
||||
from unittest import mock
|
||||
|
||||
if "websockets" not in sys.modules:
|
||||
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||
|
||||
import aibox # noqa: E402
|
||||
|
||||
|
||||
def _make_handler():
|
||||
handler = aibox.WebSocketHandler(ai_manager=mock.MagicMock())
|
||||
return handler
|
||||
|
||||
|
||||
# ==================== R1:/memos 畸形时间容错 ====================
|
||||
|
||||
class TestR1MemosCommand:
|
||||
def test_malformed_reminder_time_shows_pending(self):
|
||||
"""存量脏数据(畸形 reminder_time):/memos 按"待执行"展示,不抛 ValueError"""
|
||||
handler = _make_handler()
|
||||
try:
|
||||
handler.ai_manager.memo_db.get_all_memos.return_value = [
|
||||
{"id": 1, "title": "脏数据任务", "reminder_time": "not-a-time",
|
||||
"repeat_type": "none", "repeat_interval": 0, "is_immediate": False},
|
||||
{"id": 2, "title": "正常过期任务", "reminder_time": "2020-01-01 00:00:00",
|
||||
"repeat_type": "none", "repeat_interval": 0, "is_immediate": False},
|
||||
]
|
||||
sent = []
|
||||
|
||||
async def fake_send(ws, msg):
|
||||
sent.append(msg)
|
||||
|
||||
handler.send_message = fake_send
|
||||
ws = mock.AsyncMock()
|
||||
asyncio.run(handler.handle_command(ws, {"command": "/memos"}))
|
||||
|
||||
results = [m for m in sent if m.get("type") == "command_result"]
|
||||
errors = [m for m in sent if m.get("type") == "error"]
|
||||
assert len(results) == 1
|
||||
assert not errors # 不再落入顶层兜底"处理失败"
|
||||
content = results[0]["content"]
|
||||
assert "脏数据任务" in content and "正常过期任务" in content
|
||||
finally:
|
||||
handler.shutdown()
|
||||
|
||||
def test_valid_reminder_time_still_marks_overdue(self):
|
||||
"""合法过期时间仍标 overdue(语义不回归)"""
|
||||
handler = _make_handler()
|
||||
try:
|
||||
handler.ai_manager.memo_db.get_all_memos.return_value = [
|
||||
{"id": 1, "title": "过期任务", "reminder_time": "2020-01-01 00:00:00",
|
||||
"repeat_type": "none", "repeat_interval": 0, "is_immediate": False},
|
||||
{"id": 2, "title": "未来任务", "reminder_time": "2099-01-01 00:00:00",
|
||||
"repeat_type": "none", "repeat_interval": 0, "is_immediate": False},
|
||||
]
|
||||
sent = []
|
||||
|
||||
async def fake_send(ws, msg):
|
||||
sent.append(msg)
|
||||
|
||||
handler.send_message = fake_send
|
||||
ws = mock.AsyncMock()
|
||||
asyncio.run(handler.handle_command(ws, {"command": "/memos"}))
|
||||
|
||||
content = next(m["content"] for m in sent
|
||||
if m.get("type") == "command_result")
|
||||
overdue = aibox.i18n.get('memo_overdue')
|
||||
pending = aibox.i18n.get('memo_pending')
|
||||
overdue_line = next(l for l in content.splitlines() if "过期任务" in l)
|
||||
pending_line = next(l for l in content.splitlines() if "未来任务" in l)
|
||||
assert overdue in overdue_line and pending not in overdue_line
|
||||
assert pending in pending_line and overdue not in pending_line
|
||||
finally:
|
||||
handler.shutdown()
|
||||
|
||||
|
||||
# ==================== R2:interrupt 不再误杀 CLI 主 AI 子进程 ====================
|
||||
|
||||
class TestR2InterruptScope:
|
||||
def test_interrupt_does_not_touch_global_executor(self):
|
||||
"""WS 中断不再取消 ai_manager.command_executor(CLI 主 AI 的执行器)"""
|
||||
handler = _make_handler()
|
||||
try:
|
||||
ws = mock.AsyncMock()
|
||||
asyncio.run(handler.handle_interrupt(ws, {}))
|
||||
handler.ai_manager.command_executor.cancel.assert_not_called()
|
||||
finally:
|
||||
handler.shutdown()
|
||||
|
||||
def test_interrupt_still_cancels_session_ai(self):
|
||||
"""本会话 AI 仍被正常取消(中断能力不回归)"""
|
||||
handler = _make_handler()
|
||||
try:
|
||||
ws = mock.AsyncMock()
|
||||
ai = mock.MagicMock()
|
||||
handler.client_sessions[ws] = {"ai": ai, "is_processing": True}
|
||||
asyncio.run(handler.handle_interrupt(ws, {}))
|
||||
ai.cancel.assert_called_once()
|
||||
handler.ai_manager.command_executor.cancel.assert_not_called()
|
||||
finally:
|
||||
handler.shutdown()
|
||||
|
||||
|
||||
# ==================== R3:unregister 补齐连接回收 ====================
|
||||
|
||||
class TestR3UnregisterCleanup:
|
||||
def test_unregister_closes_websocket_and_cancels_tasks(self):
|
||||
"""broadcast 超时路径走 unregister 后:websocket 被关闭、在途任务被取消、登记清空"""
|
||||
handler = _make_handler()
|
||||
try:
|
||||
ws = mock.AsyncMock()
|
||||
|
||||
async def scenario():
|
||||
task = asyncio.create_task(asyncio.sleep(30))
|
||||
handler.connected_clients.add(ws)
|
||||
handler.client_sessions[ws] = {"ai": None, "is_processing": True}
|
||||
handler._processing_tasks[ws] = {task}
|
||||
await handler.unregister(ws)
|
||||
await asyncio.sleep(0) # 让 cancel 传播
|
||||
assert task.done() # 取消即视为完成(cancelled 或已收尾)
|
||||
|
||||
asyncio.run(scenario())
|
||||
assert ws.close.called
|
||||
assert ws not in handler.connected_clients
|
||||
assert ws not in handler.client_sessions
|
||||
assert ws not in handler._processing_tasks
|
||||
finally:
|
||||
handler.shutdown()
|
||||
|
||||
def test_unregister_idempotent_and_mock_safe(self):
|
||||
"""重复注销幂等;close 抛异常的对象不炸(AttributeError 被兜底)"""
|
||||
class _BadCloseWS:
|
||||
remote_address = ("127.0.0.1", 1)
|
||||
|
||||
async def close(self):
|
||||
raise AttributeError("no close")
|
||||
|
||||
handler = _make_handler()
|
||||
try:
|
||||
ws = _BadCloseWS()
|
||||
handler.connected_clients.add(ws)
|
||||
handler.client_sessions[ws] = {"ai": None}
|
||||
asyncio.run(handler.unregister(ws))
|
||||
asyncio.run(handler.unregister(ws)) # 二次注销不抛
|
||||
assert ws not in handler.connected_clients
|
||||
assert ws not in handler.client_sessions
|
||||
finally:
|
||||
handler.shutdown()
|
||||
@@ -0,0 +1,200 @@
|
||||
"""aibox 模型接入点管理工具(model_config.py)单元测试(2026-08-18)
|
||||
|
||||
覆盖:
|
||||
- 列表归一化(默认标记/免密/失败转移展示字段)
|
||||
- 添加:字段校验(档案名/地址/模型名/重复)、Key 两种方式、温度/超时钳制、
|
||||
失败转移引用校验、设为默认、落盘(含脱敏)
|
||||
- 删除:物理移除、fallback 悬空引用清理、默认档案标记清除
|
||||
- 连通性验证:成功 / 失败 / 档案不存在
|
||||
"""
|
||||
|
||||
import sys
|
||||
import types
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
if "websockets" not in sys.modules:
|
||||
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||
|
||||
import model_config # noqa: E402
|
||||
from common import ConfigManager # noqa: E402
|
||||
|
||||
|
||||
def _cfg(tmp_path, initial: dict = None) -> ConfigManager:
|
||||
"""构造指向临时配置文件的 ConfigManager(避免碰真实 ~/.aibox)"""
|
||||
cm = ConfigManager()
|
||||
cm._config_path = str(tmp_path / "config.json")
|
||||
cm._config = initial if initial is not None else {}
|
||||
return cm
|
||||
|
||||
|
||||
def _base_models():
|
||||
return {
|
||||
"models": {
|
||||
"deepseek": {
|
||||
"api_url": "https://api.deepseek.com/v1/chat/completions",
|
||||
"model": "deepseek-chat",
|
||||
"api_key_env": "DEEPSEEK_API_KEY",
|
||||
},
|
||||
"local": {
|
||||
"api_url": "http://localhost:11434/v1/chat/completions",
|
||||
"model": "qwen3",
|
||||
"api_key_env": "",
|
||||
},
|
||||
},
|
||||
"api": {"default_model": "deepseek"},
|
||||
}
|
||||
|
||||
|
||||
# ==================== 列表 ====================
|
||||
|
||||
class TestListProfiles:
|
||||
def test_list_normalized(self, tmp_path):
|
||||
cm = _cfg(tmp_path, _base_models())
|
||||
profiles = model_config.list_profiles(cm)
|
||||
assert len(profiles) == 2
|
||||
ds = next(p for p in profiles if p["name"] == "deepseek")
|
||||
assert ds["is_default"] is True
|
||||
assert ds["api_key_env"] == "DEEPSEEK_API_KEY"
|
||||
local = next(p for p in profiles if p["name"] == "local")
|
||||
assert local["api_key_env"] == ""
|
||||
assert local["is_default"] is False
|
||||
|
||||
def test_list_empty_models_section(self, tmp_path):
|
||||
cm = _cfg(tmp_path, {})
|
||||
assert model_config.list_profiles(cm) == []
|
||||
|
||||
def test_list_ignores_malformed_entries(self, tmp_path):
|
||||
cm = _cfg(tmp_path, {"models": {"bad": "not-a-dict"}})
|
||||
assert model_config.list_profiles(cm) == []
|
||||
|
||||
|
||||
# ==================== 校验 ====================
|
||||
|
||||
class TestValidateNewProfile:
|
||||
def test_name_rules(self, tmp_path):
|
||||
cm = _cfg(tmp_path, _base_models())
|
||||
assert model_config.validate_new_profile(cm, "", "https://x", "m")
|
||||
assert model_config.validate_new_profile(cm, "坏名字!", "https://x", "m")
|
||||
assert model_config.validate_new_profile(cm, "deepseek", "https://x", "m") # 已存在
|
||||
assert model_config.validate_new_profile(cm, "ok_name", "https://x", "m") is None
|
||||
|
||||
def test_url_scheme_and_model_required(self, tmp_path):
|
||||
cm = _cfg(tmp_path, _base_models())
|
||||
assert model_config.validate_new_profile(cm, "kimi", "ftp://x", "m")
|
||||
assert model_config.validate_new_profile(cm, "kimi", "https://x", "")
|
||||
|
||||
|
||||
# ==================== 添加 ====================
|
||||
|
||||
class TestAddProfile:
|
||||
def test_add_with_env_key_and_default(self, tmp_path):
|
||||
cm = _cfg(tmp_path, _base_models())
|
||||
msgs = model_config.add_profile(
|
||||
cm, "kimi", "https://api.moonshot.cn/v1/chat/completions",
|
||||
"moonshot-v1-8k", api_key_env="KIMI_KEY",
|
||||
display_name="Kimi", set_default=True)
|
||||
assert any("已添加" in m for m in msgs)
|
||||
assert any("设为默认" in m for m in msgs)
|
||||
assert cm.get("api.default_model") == "kimi"
|
||||
assert cm.get("models.kimi.api_key_env") == "KIMI_KEY"
|
||||
assert cm.get("models.kimi.api_url").startswith("https://")
|
||||
# 落盘文件存在且可解析
|
||||
import json
|
||||
disk = json.loads((tmp_path / "config.json").read_text(encoding="utf-8"))
|
||||
assert disk["models"]["kimi"]["model"] == "moonshot-v1-8k"
|
||||
|
||||
def test_add_without_key_is_exempt_endpoint(self, tmp_path):
|
||||
cm = _cfg(tmp_path, {})
|
||||
msgs = model_config.add_profile(
|
||||
cm, "local", "http://localhost:11434/v1/chat/completions",
|
||||
"qwen3", api_key_env="")
|
||||
assert any("免密" in m for m in msgs)
|
||||
assert cm.get("models.local.api_key_env") == ""
|
||||
|
||||
def test_add_rejects_invalid_fallback(self, tmp_path):
|
||||
cm = _cfg(tmp_path, _base_models())
|
||||
msgs = model_config.add_profile(
|
||||
cm, "kimi", "https://x/v1", "m1",
|
||||
fallback=["ghost", "deepseek"])
|
||||
assert any("失败转移" in m for m in msgs)
|
||||
assert cm.get("models.kimi") is None # 未写入
|
||||
|
||||
def test_add_temperature_timeout_clamped(self, tmp_path):
|
||||
cm = _cfg(tmp_path, {})
|
||||
model_config.add_profile(cm, "a", "https://x/v1", "m1",
|
||||
temperature=9.9, timeout=1)
|
||||
assert cm.get("models.a.temperature") == 2.0
|
||||
assert cm.get("models.a.timeout") == 5
|
||||
|
||||
def test_add_invalid_temperature_rejected(self, tmp_path):
|
||||
cm = _cfg(tmp_path, {})
|
||||
msgs = model_config.add_profile(cm, "a", "https://x/v1", "m1",
|
||||
temperature="abc")
|
||||
assert any("温度" in m for m in msgs)
|
||||
assert cm.get("models.a") is None
|
||||
|
||||
|
||||
# ==================== 删除 ====================
|
||||
|
||||
class TestDeleteProfile:
|
||||
def test_delete_cleans_fallback_refs(self, tmp_path):
|
||||
cm = _cfg(tmp_path, {
|
||||
"models": {
|
||||
"main": {"api_url": "https://a", "model": "m1",
|
||||
"fallback": ["backup"]},
|
||||
"backup": {"api_url": "https://b", "model": "m2"},
|
||||
},
|
||||
})
|
||||
msgs = model_config.delete_profile(cm, "backup")
|
||||
assert any("已删除" in m for m in msgs)
|
||||
assert any("失败转移引用" in m for m in msgs)
|
||||
assert cm.get("models.main.fallback") == []
|
||||
|
||||
def test_delete_default_clears_default_flag(self, tmp_path):
|
||||
cm = _cfg(tmp_path, _base_models())
|
||||
model_config.delete_profile(cm, "deepseek")
|
||||
assert cm.get("api.default_model") is None
|
||||
|
||||
def test_delete_nonexistent(self, tmp_path):
|
||||
cm = _cfg(tmp_path, _base_models())
|
||||
msgs = model_config.delete_profile(cm, "ghost")
|
||||
assert any("不存在" in m for m in msgs)
|
||||
|
||||
|
||||
# ==================== 连通性验证 ====================
|
||||
|
||||
class TestTestProfile:
|
||||
def test_connectivity_ok(self, tmp_path):
|
||||
cm = _cfg(tmp_path, _base_models())
|
||||
fake = mock.MagicMock()
|
||||
fake.chat_completion.return_value = "pong"
|
||||
with mock.patch("llm_client.create_llm_client", return_value=fake):
|
||||
r = model_config.test_profile(cm, "deepseek")
|
||||
assert "连通正常" in r and "pong" in r
|
||||
|
||||
def test_connectivity_failure(self, tmp_path):
|
||||
cm = _cfg(tmp_path, _base_models())
|
||||
fake = mock.MagicMock()
|
||||
fake.chat_completion.side_effect = RuntimeError("连接被拒绝")
|
||||
with mock.patch("llm_client.create_llm_client", return_value=fake):
|
||||
r = model_config.test_profile(cm, "deepseek")
|
||||
assert "验证失败" in r and "连接被拒绝" in r
|
||||
|
||||
def test_connectivity_missing_profile(self, tmp_path):
|
||||
cm = _cfg(tmp_path, _base_models())
|
||||
assert "不存在" in model_config.test_profile(cm, "ghost")
|
||||
|
||||
|
||||
# ==================== 落盘失败感知 ====================
|
||||
|
||||
class TestSaveFailure:
|
||||
def test_add_reports_save_failure(self, tmp_path):
|
||||
cm = _cfg(tmp_path, {})
|
||||
# J5:写路径改为原子写,磁盘失败经原子写函数模拟
|
||||
with mock.patch("common.config.atomic_write_text",
|
||||
side_effect=PermissionError("disk ro")):
|
||||
msgs = model_config.add_profile(
|
||||
cm, "a", "https://x/v1", "m1", api_key_env="K")
|
||||
assert any("保存失败" in m for m in msgs)
|
||||
@@ -820,8 +820,9 @@ class TestL4aChatCompletionFailover:
|
||||
assert fc.chat_completion([{"role": "user", "content": "q2"}]) == "r2"
|
||||
assert mp.call_args_list[0].args[0].startswith("https://api.moonshot.cn")
|
||||
# 冷却过期(手动回拨失败时间戳):恢复走主档案
|
||||
fc._fail_at["deepseek"] = (_time.monotonic() - fc.cooldown_seconds - 1,
|
||||
fc.cooldown_seconds)
|
||||
# T6:冷却键改为客户端实例 id,种冷却态需走 _cool_key
|
||||
fc._fail_at[fc._cool_key(fc._clients[0])] = (_time.monotonic() - fc.cooldown_seconds - 1,
|
||||
fc.cooldown_seconds)
|
||||
with mock.patch("llm_client.requests.post",
|
||||
return_value=_resp(200, content="r3")) as mp:
|
||||
assert fc.chat_completion([{"role": "user", "content": "q3"}]) == "r3"
|
||||
@@ -998,8 +999,8 @@ class TestL4bStreamFailover:
|
||||
primary = self._client("deepseek", stream=_fake_stream(events=self.EVENTS))
|
||||
fallback = self._client("kimi", stream=_fake_stream(events=self.EVENTS))
|
||||
fc = self._fc(primary, fallback)
|
||||
fc._fail_at = {"deepseek": (_time.monotonic(), fc.cooldown_seconds),
|
||||
"kimi": (_time.monotonic(), fc.cooldown_seconds)}
|
||||
fc._fail_at = {fc._cool_key(primary): (_time.monotonic(), fc.cooldown_seconds),
|
||||
fc._cool_key(fallback): (_time.monotonic(), fc.cooldown_seconds)}
|
||||
with pytest.raises(RuntimeError, match="冷却期"):
|
||||
list(fc.stream_chat([{"role": "user", "content": "q"}]))
|
||||
primary.stream_chat.assert_not_called()
|
||||
@@ -1010,7 +1011,7 @@ class TestL4bStreamFailover:
|
||||
primary = self._client("deepseek", stream=_fake_stream(events=self.EVENTS))
|
||||
fallback = self._client("kimi", stream=_fake_stream(events=self.EVENTS))
|
||||
fc = self._fc(primary, fallback)
|
||||
fc._fail_at["deepseek"] = (_time.monotonic(), fc.cooldown_seconds)
|
||||
fc._fail_at[fc._cool_key(primary)] = (_time.monotonic(), fc.cooldown_seconds)
|
||||
assert list(fc.stream_chat([{"role": "user", "content": "q"}])) == self.EVENTS
|
||||
primary.stream_chat.assert_not_called()
|
||||
|
||||
@@ -1102,7 +1103,7 @@ class TestL4cObservabilityAndBudget:
|
||||
stream=_fake_stream(events=self.EVENTS))
|
||||
fallback = self._mk_client("kimi", stream=_fake_stream(events=self.EVENTS))
|
||||
fc = self._fc(primary, fallback)
|
||||
fc._fail_at["deepseek"] = (_time.monotonic(), fc.cooldown_seconds)
|
||||
fc._fail_at[fc._cool_key(primary)] = (_time.monotonic(), fc.cooldown_seconds)
|
||||
assert list(fc.stream_chat([{"role": "user", "content": "q"}])) == self.EVENTS
|
||||
s = fc.get_stats()
|
||||
assert s == {"requests": 1, "failovers": 0, "cooldown_skips": 1, "exhausted": 0}
|
||||
@@ -1231,13 +1232,14 @@ class TestL4cObservabilityAndBudget:
|
||||
# 第一次:主失败(retry_after=300)→ fallback
|
||||
assert fc.chat_completion([{"role": "user", "content": "q1"}]) == "from-kimi"
|
||||
# 冷却中(默认 60s 已过但 300s 未到):仍走 fallback
|
||||
t, _ = fc._fail_at["deepseek"]
|
||||
fc._fail_at["deepseek"] = (t - 100, 300)
|
||||
_key = fc._cool_key(primary)
|
||||
t, _ = fc._fail_at[_key]
|
||||
fc._fail_at[_key] = (t - 100, 300)
|
||||
primary.chat_completion = primary_ok
|
||||
assert fc.chat_completion([{"role": "user", "content": "q2"}]) == "from-kimi"
|
||||
primary_ok.assert_not_called()
|
||||
# 300s 冷却过期:恢复主档案
|
||||
fc._fail_at["deepseek"] = (t - 301, 300)
|
||||
fc._fail_at[_key] = (t - 301, 300)
|
||||
assert fc.chat_completion([{"role": "user", "content": "q3"}]) == "from-ds"
|
||||
primary_ok.assert_called_once()
|
||||
|
||||
|
||||
@@ -0,0 +1,829 @@
|
||||
"""tests.test_n_fixes - Batch N 修复单元测试
|
||||
|
||||
波次 3 平台兼容(2026-08-21,审计第八轮 T 系列):
|
||||
- T3 basic_tools 平台感知引用 + Windows PowerShell/curl.exe 指引
|
||||
- T11 Windows 进程树收割(taskkill /T /F,executors + process_manager)
|
||||
- T25 MemoryDatabase 统一连接工厂(busy_timeout)
|
||||
- T26 SkillManager 原子 reload(中途不清空 self.skills)
|
||||
- T27 atomic_write_text 对 Windows PermissionError 短退避重试
|
||||
|
||||
波次 4 LLM 健壮性(2026-08-21):
|
||||
- T10 调度器 wait_for 硬超时兜底
|
||||
- T19 FailoverClient 异常收窄(永久配置错误不转移不冷却)
|
||||
- T20 流式响应未收到 [DONE] 抛 TransientLLMError(不再静默截断)
|
||||
- T21 env 覆盖黑名单(LD_*/DYLD_*/PATH/PYTHONPATH 等劫持面)
|
||||
- T23 压缩后 tool 配对双向修复(assistant.tool_calls 失配剥离)
|
||||
(T14 FTS5 经复现证伪:JOIN 别名场景下表名 MATCH 正常命中,
|
||||
全链路中文检索长词走 trigram、短词走 LIKE 兜底,无需修复)
|
||||
|
||||
波次 5 权限隔离(2026-08-21,产品决策:T1 仅告警不加鉴权、T2 维持现状):
|
||||
- T22 定时任务精简 profile(task:无 update_identity/update_soul/set_proxy)
|
||||
+ ProcessManager owner 隔离(写/杀权限,观察面共享)
|
||||
- T-L17 model_config 按域名推荐 env 名(防 DeepSeek 凭据发往他厂)
|
||||
- T-L18 list_profiles 输出过滤控制字符(防手写 config ANSI 注入)
|
||||
|
||||
Windows 分支在 Linux 上以 mock.patch("os.name", "nt") 模拟验证命令生成
|
||||
逻辑(不实际执行生成的命令)。
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sqlite3
|
||||
import threading
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from common.atomic import atomic_write_text
|
||||
from subsystems import MemoryDatabase, SkillManager
|
||||
from subsystems.executors import SubprocessCommandExecutor, filter_env_overrides
|
||||
from tools.basic_tools import (
|
||||
ListFilesTool,
|
||||
ReadFileTool,
|
||||
SimpleWebCrawlerTool,
|
||||
WriteFileTool,
|
||||
_quote_arg,
|
||||
)
|
||||
|
||||
|
||||
# ==================== T3:平台感知引用 ====================
|
||||
|
||||
class TestQuoteArg:
|
||||
def test_posix_uses_shlex(self):
|
||||
quoted, err = _quote_arg("/tmp/my file.txt")
|
||||
assert err is None
|
||||
assert quoted == "'/tmp/my file.txt'"
|
||||
|
||||
@pytest.mark.skipif(os.name != "posix", reason="POSIX 引用语义")
|
||||
def test_posix_specials_escaped(self):
|
||||
quoted, err = _quote_arg("a`b$c$(d)\\e")
|
||||
assert err is None
|
||||
assert quoted == "'a`b$c$(d)\\e'"
|
||||
|
||||
def test_win_wraps_double_quotes(self):
|
||||
with mock.patch("os.name", "nt"):
|
||||
quoted, err = _quote_arg(r"C:\Program Files\my app\doc.txt")
|
||||
assert err is None
|
||||
assert quoted == r'"C:\Program Files\my app\doc.txt"'
|
||||
|
||||
def test_win_rejects_double_quote(self):
|
||||
with mock.patch("os.name", "nt"):
|
||||
quoted, err = _quote_arg('evil" & del /f /s /q C:\\')
|
||||
assert quoted == ""
|
||||
assert err is not None and "双引号" in err
|
||||
|
||||
def test_win_rejects_cmd_var_expansion(self):
|
||||
with mock.patch("os.name", "nt"):
|
||||
quoted, err = _quote_arg("https://evil.com/?x=%PATH%")
|
||||
assert quoted == ""
|
||||
assert err is not None and "注入" in err
|
||||
|
||||
def test_win_allows_url_percent_encoding(self):
|
||||
"""URL 编码(%20/%E4 等)非 %变量名% 形态,不拒绝。"""
|
||||
with mock.patch("os.name", "nt"):
|
||||
quoted, err = _quote_arg("https://example.com/a%20b?x=1&y=2")
|
||||
assert err is None
|
||||
assert quoted == '"https://example.com/a%20b?x=1&y=2"'
|
||||
|
||||
|
||||
# ==================== T3:Windows 指引命令生成 ====================
|
||||
|
||||
class TestReadFileWindows:
|
||||
def test_head_powershell(self, tmp_path):
|
||||
f = tmp_path / "data.txt"
|
||||
f.write_text("hello\n" * 30, encoding="utf-8")
|
||||
with mock.patch("os.name", "nt"):
|
||||
r = ReadFileTool().execute(filepath=str(f), head=True, max_lines=10)
|
||||
assert "powershell" in r and "Get-Content" in r
|
||||
assert "run_command" in r
|
||||
assert "sed -n" not in r and "head -n" not in r
|
||||
|
||||
def test_tail_powershell(self, tmp_path):
|
||||
f = tmp_path / "data.txt"
|
||||
f.write_text("hello\n" * 30, encoding="utf-8")
|
||||
with mock.patch("os.name", "nt"):
|
||||
r = ReadFileTool().execute(filepath=str(f), tail=True, max_lines=5)
|
||||
assert "Get-Content" in r and "-Tail 5" in r
|
||||
|
||||
def test_range_powershell(self, tmp_path):
|
||||
f = tmp_path / "data.txt"
|
||||
f.write_text("hello\n" * 30, encoding="utf-8")
|
||||
with mock.patch("os.name", "nt"):
|
||||
r = ReadFileTool().execute(filepath=str(f), start_line=5, end_line=15)
|
||||
assert "Select-Object -Skip 4 -First 11" in r
|
||||
|
||||
def test_posix_unchanged(self, tmp_path):
|
||||
f = tmp_path / "data.txt"
|
||||
f.write_text("hello\n" * 30, encoding="utf-8")
|
||||
r = ReadFileTool().execute(filepath=str(f), start_line=1, end_line=5)
|
||||
assert "sed -n '1,5p'" in r
|
||||
|
||||
|
||||
class TestWriteFileWindows:
|
||||
def test_base64_set_content(self, tmp_path):
|
||||
target = tmp_path / "sub" / "out.txt"
|
||||
content = "line1\nline2\n含中文 \"quoted\" & special"
|
||||
with mock.patch("os.name", "nt"):
|
||||
r = WriteFileTool().execute(filepath=str(target), content=content, mode="w")
|
||||
assert "Set-Content" in r
|
||||
assert "FromBase64String" in r
|
||||
# 明文内容(含元字符)不得出现在命令行中
|
||||
assert "quoted" not in r and "special" not in r
|
||||
|
||||
def test_append_uses_add_content(self, tmp_path):
|
||||
with mock.patch("os.name", "nt"):
|
||||
r = WriteFileTool().execute(
|
||||
filepath=str(tmp_path / "a.txt"), content="more", mode="a")
|
||||
assert "Add-Content" in r
|
||||
|
||||
def test_oversize_rejected(self, tmp_path):
|
||||
with mock.patch("os.name", "nt"):
|
||||
r = WriteFileTool().execute(
|
||||
filepath=str(tmp_path / "big.txt"), content="x" * 10000, mode="w")
|
||||
assert "过大" in r and "追加" in r
|
||||
|
||||
|
||||
class TestListFilesWindows:
|
||||
def test_get_child_item(self, tmp_path):
|
||||
with mock.patch("os.name", "nt"):
|
||||
r = ListFilesTool().execute(directory=str(tmp_path), show_hidden=True)
|
||||
assert "Get-ChildItem" in r and "-Force" in r
|
||||
assert "ls -l" not in r
|
||||
|
||||
def test_posix_unchanged(self, tmp_path):
|
||||
r = ListFilesTool().execute(directory=str(tmp_path))
|
||||
assert "ls -l " in r
|
||||
|
||||
|
||||
class TestWebWindows:
|
||||
def test_get_curl_exe_amp_url_quoted(self):
|
||||
tool = SimpleWebCrawlerTool()
|
||||
url = "https://example.com/?a=1&b=2"
|
||||
with mock.patch("os.name", "nt"):
|
||||
r = tool.execute(action="get", url=url)
|
||||
assert "curl.exe" in r
|
||||
# & 必须在双引号内(cmd 双引号内为字面量)
|
||||
assert '"https://example.com/?a=1&b=2"' in r
|
||||
|
||||
def test_cmd_var_url_rejected(self):
|
||||
tool = SimpleWebCrawlerTool()
|
||||
with mock.patch("os.name", "nt"):
|
||||
r = tool.execute(action="get", url="https://evil.com/?x=%PATH%")
|
||||
assert "拒绝" in r
|
||||
|
||||
def test_text_action_powershell_pipe(self):
|
||||
tool = SimpleWebCrawlerTool()
|
||||
with mock.patch("os.name", "nt"):
|
||||
r = tool.execute(action="text", url="https://example.com/")
|
||||
assert "curl.exe" in r and "powershell" in r
|
||||
|
||||
def test_posix_unchanged(self):
|
||||
tool = SimpleWebCrawlerTool()
|
||||
r = tool.execute(action="get", url="https://example.com/")
|
||||
assert "curl -sL" in r
|
||||
|
||||
|
||||
# ==================== T11:Windows 进程树收割 ====================
|
||||
|
||||
class TestTerminateTreeWindows:
|
||||
def test_executors_taskkill_tree(self):
|
||||
ex = SubprocessCommandExecutor()
|
||||
proc = mock.Mock()
|
||||
proc.pid = 12345
|
||||
proc.poll.return_value = None
|
||||
with mock.patch("os.name", "nt"), \
|
||||
mock.patch("subprocess.run") as run:
|
||||
ex._terminate_tree(proc, force=True)
|
||||
run.assert_called_once()
|
||||
assert run.call_args[0][0] == ["taskkill", "/PID", "12345", "/T", "/F"]
|
||||
proc.kill.assert_called_once() # 兜底单进程 kill
|
||||
|
||||
def test_executors_taskkill_process_gone(self):
|
||||
"""taskkill 目标已退出:兜底 kill 不执行(poll 非 None)。"""
|
||||
ex = SubprocessCommandExecutor()
|
||||
proc = mock.Mock()
|
||||
proc.pid = 12346
|
||||
proc.poll.return_value = 0
|
||||
with mock.patch("os.name", "nt"), \
|
||||
mock.patch("subprocess.run") as run:
|
||||
ex._terminate_tree(proc, force=False)
|
||||
run.assert_called_once()
|
||||
proc.kill.assert_not_called()
|
||||
|
||||
def test_process_manager_taskkill_tree(self):
|
||||
from subsystems.process_manager import ProcessManager
|
||||
pm = ProcessManager()
|
||||
entry = mock.Mock()
|
||||
entry.popen.pid = 999
|
||||
entry.popen.poll.return_value = None
|
||||
with mock.patch("os.name", "nt"), \
|
||||
mock.patch("subprocess.run") as run:
|
||||
pm._terminate_tree(entry, force=True)
|
||||
assert run.call_args[0][0] == ["taskkill", "/PID", "999", "/T", "/F"]
|
||||
|
||||
@pytest.mark.skipif(os.name != "posix", reason="POSIX 进程组语义")
|
||||
def test_posix_no_taskkill(self):
|
||||
"""POSIX 分支不调用 taskkill(走 killpg)。"""
|
||||
import subprocess as _sp
|
||||
ex = SubprocessCommandExecutor()
|
||||
proc = _sp.Popen(["sleep", "30"], start_new_session=True)
|
||||
try:
|
||||
with mock.patch("subprocess.run") as run:
|
||||
ex._terminate_tree(proc, force=False)
|
||||
run.assert_not_called()
|
||||
finally:
|
||||
try:
|
||||
import signal as _sig
|
||||
os.killpg(os.getpgid(proc.pid), _sig.SIGKILL)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# ==================== T25:MemoryDatabase busy_timeout ====================
|
||||
|
||||
class TestMemoryDbBusyTimeout:
|
||||
def test_connect_sets_busy_timeout(self, tmp_path):
|
||||
db = MemoryDatabase(db_path=str(tmp_path / "mem.db"))
|
||||
conn = db._connect()
|
||||
try:
|
||||
value = conn.execute("PRAGMA busy_timeout").fetchone()[0]
|
||||
finally:
|
||||
conn.close()
|
||||
assert value == 10000
|
||||
|
||||
def test_all_paths_via_factory(self, tmp_path):
|
||||
"""写/读路径不再直连 sqlite3.connect(统一走 _connect 工厂)。"""
|
||||
db = MemoryDatabase(db_path=str(tmp_path / "mem.db"))
|
||||
mid = db.add_memory(topic="t", summary="s", content="c",
|
||||
start_time="2026-01-01 00:00:00",
|
||||
end_time="2026-01-01 00:01:00")
|
||||
assert mid is not None
|
||||
rows = db.get_recent_memories(limit=5)
|
||||
assert any(r["id"] == mid for r in rows)
|
||||
|
||||
|
||||
# ==================== T26:SkillManager 原子 reload ====================
|
||||
|
||||
@pytest.fixture
|
||||
def isolated_skill_dir(tmp_path, monkeypatch):
|
||||
from common import constants
|
||||
skills_dir = tmp_path / "skills"
|
||||
skills_dir.mkdir()
|
||||
monkeypatch.setattr(type(constants), "SKILLS_DIR",
|
||||
property(lambda self: str(skills_dir)))
|
||||
return skills_dir
|
||||
|
||||
|
||||
class TestSkillReloadAtomic:
|
||||
def test_skills_visible_during_reload(self, isolated_skill_dir):
|
||||
"""reload 过程中 self.skills 保持旧表(T26:不再先清空)。"""
|
||||
(isolated_skill_dir / "a.md").write_text("# A\n内容A", encoding="utf-8")
|
||||
sm = SkillManager()
|
||||
assert sm.count() == 1
|
||||
|
||||
(isolated_skill_dir / "b.md").write_text("# B\n内容B", encoding="utf-8")
|
||||
original = sm._read_markdown_skill
|
||||
observed = []
|
||||
calls = []
|
||||
|
||||
def spy(path):
|
||||
calls.append(path)
|
||||
if len(calls) == 2:
|
||||
observed.append(dict(sm.skills))
|
||||
return original(path)
|
||||
|
||||
sm._read_markdown_skill = spy
|
||||
sm.reload()
|
||||
|
||||
# 加载第 2 个技能时,旧表完整可见(未被清空/半截)
|
||||
assert len(observed) == 1
|
||||
assert set(observed[0].keys()) == {"a"}
|
||||
# reload 完成后一次性替换为新表
|
||||
assert set(sm.skills.keys()) == {"a", "b"}
|
||||
|
||||
def test_reload_missing_dir_keeps_old(self, isolated_skill_dir, monkeypatch):
|
||||
"""目录消失时清空(旧行为保留),正常 reload 不经过空表中间态。"""
|
||||
from common import constants
|
||||
(isolated_skill_dir / "a.md").write_text("# A\n内容A", encoding="utf-8")
|
||||
sm = SkillManager()
|
||||
assert sm.count() == 1
|
||||
monkeypatch.setattr(type(constants), "SKILLS_DIR",
|
||||
property(lambda self: str(isolated_skill_dir / "nope")))
|
||||
sm.reload()
|
||||
assert sm.count() == 0
|
||||
|
||||
|
||||
# ==================== T27:atomic_write_text 重试 ====================
|
||||
|
||||
class TestAtomicWriteRetry:
|
||||
def test_retries_then_succeeds(self, tmp_path):
|
||||
target = tmp_path / "cfg.json"
|
||||
real_replace = os.replace
|
||||
calls = []
|
||||
|
||||
def flaky_replace(src, dst):
|
||||
calls.append((src, dst))
|
||||
if len(calls) < 3:
|
||||
raise PermissionError(5, "Access is denied")
|
||||
return real_replace(src, dst)
|
||||
|
||||
with mock.patch("os.replace", side_effect=flaky_replace):
|
||||
atomic_write_text(str(target), '{"k": 1}')
|
||||
assert len(calls) == 3
|
||||
assert target.read_text(encoding="utf-8") == '{"k": 1}'
|
||||
|
||||
def test_raises_after_exhausted(self, tmp_path):
|
||||
target = tmp_path / "cfg.json"
|
||||
with mock.patch("os.replace",
|
||||
side_effect=PermissionError(5, "Access is denied")):
|
||||
with pytest.raises(PermissionError):
|
||||
atomic_write_text(str(target), "x")
|
||||
assert not target.exists()
|
||||
# 临时文件已清理
|
||||
leftovers = [p for p in tmp_path.iterdir() if p.name.startswith(".aibox_tmp_")]
|
||||
assert leftovers == []
|
||||
|
||||
def test_normal_write_no_retry(self, tmp_path):
|
||||
target = tmp_path / "ok.txt"
|
||||
with mock.patch("os.replace", wraps=os.replace) as spy:
|
||||
atomic_write_text(str(target), "hello")
|
||||
spy.assert_called_once()
|
||||
assert target.read_text(encoding="utf-8") == "hello"
|
||||
|
||||
|
||||
# ==================== T19:FailoverClient 异常收窄 ====================
|
||||
|
||||
def _make_failover_client():
|
||||
"""deepseek(→kimi fallback) 链,Key 走 mock 环境变量。"""
|
||||
from llm_client import create_llm_client
|
||||
from test_models_profile import MODELS, _cfg
|
||||
models = {
|
||||
"deepseek": dict(MODELS["deepseek"], fallback=["kimi"]),
|
||||
"kimi": dict(MODELS["kimi"]),
|
||||
}
|
||||
env = {"DEEPSEEK_API_KEY": "sk-ds", "MOONSHOT_API_KEY": "sk-moon"}
|
||||
with mock.patch.dict(os.environ, env):
|
||||
return create_llm_client(_cfg(models=models))
|
||||
|
||||
|
||||
class TestT19TransientExceptionScope:
|
||||
def test_permanent_schema_error_not_failed_over(self):
|
||||
"""MissingSchema(api_url 手写错等永久配置错误)→ 流式路径直接上抛,
|
||||
不转移(只打一次请求)、不计 failover、不冷却。"""
|
||||
fc = _make_failover_client()
|
||||
with mock.patch("llm_client.requests.post",
|
||||
side_effect=requests.exceptions.MissingSchema("bad url")) as mp:
|
||||
with pytest.raises(requests.exceptions.MissingSchema):
|
||||
list(fc.stream_chat([{"role": "user", "content": "hi"}]))
|
||||
assert len(mp.call_args_list) == 1
|
||||
assert fc.get_stats()["failovers"] == 0
|
||||
|
||||
def test_permanent_error_no_cooldown_pollution(self):
|
||||
"""永久错误后档案不进冷却——下次调用仍打主档案。"""
|
||||
fc = _make_failover_client()
|
||||
ok_stream = _stream_response([
|
||||
'data: {"choices":[{"delta":{"content":"recovered"}}]}',
|
||||
'data: [DONE]'])
|
||||
with mock.patch("llm_client.requests.post",
|
||||
side_effect=[requests.exceptions.MissingSchema("bad"),
|
||||
ok_stream]) as mp:
|
||||
with pytest.raises(requests.exceptions.MissingSchema):
|
||||
list(fc.stream_chat([{"role": "user", "content": "hi"}]))
|
||||
events = list(fc.stream_chat([{"role": "user", "content": "hi"}]))
|
||||
# 第二次仍先打主档案(deepseek),未进冷却跳过
|
||||
assert mp.call_args_list[1].args[0].startswith("https://api.deepseek.com")
|
||||
assert any(e["type"] == "content" and e["text"] == "recovered" for e in events)
|
||||
assert fc.get_stats()["cooldown_skips"] == 0
|
||||
|
||||
def test_connection_error_still_fails_over(self):
|
||||
"""瞬时连接错误仍正常转移(收窄不误伤)。"""
|
||||
fc = _make_failover_client()
|
||||
from test_models_profile import _resp
|
||||
with mock.patch("llm_client.requests.post",
|
||||
side_effect=[requests.exceptions.ConnectionError("boom"),
|
||||
_resp(200, content="ok-conn")]):
|
||||
assert fc.chat_completion([{"role": "user", "content": "hi"}]) == "ok-conn"
|
||||
|
||||
def test_chunked_encoding_error_fails_over(self):
|
||||
"""ChunkedEncodingError(响应体中途损坏)属瞬时故障 → 转移。"""
|
||||
fc = _make_failover_client()
|
||||
from test_models_profile import _resp
|
||||
with mock.patch("llm_client.requests.post",
|
||||
side_effect=[requests.exceptions.ChunkedEncodingError("broken"),
|
||||
_resp(200, content="ok-chunked")]):
|
||||
assert fc.chat_completion([{"role": "user", "content": "hi"}]) == "ok-chunked"
|
||||
|
||||
|
||||
# ==================== T20:不完整流信号 ====================
|
||||
|
||||
def _stream_response(lines):
|
||||
m = mock.MagicMock()
|
||||
m.status_code = 200
|
||||
m.iter_lines.return_value = [l.encode("utf-8") for l in lines]
|
||||
return m
|
||||
|
||||
|
||||
def _make_deepseek_client():
|
||||
from llm_client import DeepSeekClient
|
||||
return DeepSeekClient(
|
||||
api_key="sk-test", model="test-model",
|
||||
base_url="https://api.example/v1/chat/completions")
|
||||
|
||||
|
||||
class TestT20IncompleteStream:
|
||||
def test_missing_done_raises_transient(self):
|
||||
"""流中途断开(iter_lines 耗尽无 [DONE])→ TransientLLMError,
|
||||
不再静默视为成功导致回答无声截断。"""
|
||||
from llm_client import TransientLLMError
|
||||
c = _make_deepseek_client()
|
||||
lines = [
|
||||
'data: {"choices":[{"delta":{"content":"he"}}]}',
|
||||
'data: {"choices":[{"delta":{"content":"llo"}}]}',
|
||||
# 连接断开:无 [DONE],末行可能还是半截 JSON
|
||||
'data: {"choices":[{"delta":{"cont',
|
||||
]
|
||||
with mock.patch("llm_client.requests.post",
|
||||
return_value=_stream_response(lines)):
|
||||
with pytest.raises(TransientLLMError, match="不完整"):
|
||||
list(c.stream_chat([{"role": "user", "content": "hi"}]))
|
||||
|
||||
def test_with_done_completes_normally(self):
|
||||
from llm_client import TransientLLMError
|
||||
c = _make_deepseek_client()
|
||||
lines = [
|
||||
'data: {"choices":[{"delta":{"content":"hi"}}]}',
|
||||
'data: [DONE]',
|
||||
]
|
||||
with mock.patch("llm_client.requests.post",
|
||||
return_value=_stream_response(lines)):
|
||||
events = list(c.stream_chat([{"role": "user", "content": "hi"}]))
|
||||
assert events[-1] == {"type": "done"}
|
||||
assert any(e["type"] == "content" and e["text"] == "hi" for e in events)
|
||||
|
||||
def test_failover_first_event_incomplete_stream(self):
|
||||
"""FailoverClient:首事件前不完整流(0 事件 + 无 [DONE])→ 可转移。"""
|
||||
fc = _make_failover_client()
|
||||
# 主档案:断流无任何事件;fallback(kimi):正常完整流
|
||||
broken = _stream_response(['data: {"choices":[{"delta'])
|
||||
ok_stream = _stream_response([
|
||||
'data: {"choices":[{"delta":{"content":"from-kimi"}}]}',
|
||||
'data: [DONE]'])
|
||||
with mock.patch("llm_client.requests.post",
|
||||
side_effect=[broken, ok_stream]):
|
||||
events = list(fc.stream_chat([{"role": "user", "content": "hi"}]))
|
||||
assert any(e["type"] == "content" and e["text"] == "from-kimi" for e in events)
|
||||
assert events[-1]["type"] == "done"
|
||||
|
||||
|
||||
# ==================== T21:env 覆盖黑名单 ====================
|
||||
|
||||
class TestT21EnvDenylist:
|
||||
def test_denies_hijack_keys(self):
|
||||
safe, rejected = filter_env_overrides({
|
||||
"LD_PRELOAD": "/evil", "PATH": "/bin", "PYTHONPATH": "/p",
|
||||
"MY_VAR": "ok"})
|
||||
assert safe == {"MY_VAR": "ok"}
|
||||
assert set(rejected) == {"LD_PRELOAD", "PATH", "PYTHONPATH"}
|
||||
|
||||
def test_case_insensitive_and_prefixes(self):
|
||||
safe, rejected = filter_env_overrides({
|
||||
"ld_preload": "x", "DYLD_INSERT_LIBRARIES": "y", "path": "z",
|
||||
"PSMODULEPATH": "w", "OK1": "1"})
|
||||
assert safe == {"OK1": "1"}
|
||||
assert len(rejected) == 4
|
||||
|
||||
def test_empty_and_none(self):
|
||||
assert filter_env_overrides(None) == ({}, [])
|
||||
assert filter_env_overrides({}) == ({}, [])
|
||||
|
||||
def test_executor_blocks_denied_env(self):
|
||||
"""执行路径集成:正常 env 注入生效,黑名单键被拒。"""
|
||||
ex = SubprocessCommandExecutor()
|
||||
r = ex.execute(
|
||||
"echo $T21VAR",
|
||||
env={"T21VAR": "visible", "LD_PRELOAD": "/nonexistent/libevil.so"})
|
||||
assert "visible" in r
|
||||
# LD_PRELOAD 未注入(若注入且文件不存在,loader 报错会进合并输出)
|
||||
assert "/nonexistent" not in r
|
||||
|
||||
|
||||
# ==================== T23:压缩后 tool 配对双向修复 ====================
|
||||
|
||||
def _make_compressor():
|
||||
from subsystems.conversation_compressor import ConversationCompressor
|
||||
return ConversationCompressor(llm_client=None)
|
||||
|
||||
|
||||
class TestT23ToolPairValidation:
|
||||
def test_orphan_tool_calls_stripped(self):
|
||||
"""assistant.tool_calls 的 tool 结果被压缩掉 → 剥离 tool_calls。"""
|
||||
comp = _make_compressor()
|
||||
msgs = [
|
||||
{"role": "user", "content": "q"},
|
||||
{"role": "assistant", "content": "看结果",
|
||||
"tool_calls": [{"id": "t1", "type": "function",
|
||||
"function": {"name": "read", "arguments": "{}"}}]},
|
||||
{"role": "user", "content": "next"},
|
||||
]
|
||||
out = comp._validate_message_structure(msgs)
|
||||
assert len(out) == 3
|
||||
assert "tool_calls" not in out[1]
|
||||
assert out[1]["content"] == "看结果" # 原有内容保留
|
||||
|
||||
def test_empty_content_filled_when_stripped(self):
|
||||
comp = _make_compressor()
|
||||
msgs = [
|
||||
{"role": "assistant", "content": None,
|
||||
"tool_calls": [{"id": "t1", "type": "function",
|
||||
"function": {"name": "read", "arguments": "{}"}}]},
|
||||
]
|
||||
out = comp._validate_message_structure(msgs)
|
||||
assert "tool_calls" not in out[0]
|
||||
assert out[0]["content"] == "" # 防产生非法空消息
|
||||
|
||||
def test_partial_pair_keeps_matched_calls(self):
|
||||
"""两个 tool_calls 只剩一个配对 → 保留配对的那个。"""
|
||||
comp = _make_compressor()
|
||||
msgs = [
|
||||
{"role": "assistant", "content": "run",
|
||||
"tool_calls": [
|
||||
{"id": "t1", "type": "function",
|
||||
"function": {"name": "a", "arguments": "{}"}},
|
||||
{"id": "t2", "type": "function",
|
||||
"function": {"name": "b", "arguments": "{}"}}]},
|
||||
{"role": "tool", "tool_call_id": "t2", "content": "result-b"},
|
||||
]
|
||||
out = comp._validate_message_structure(msgs)
|
||||
assert len(out) == 2
|
||||
assert [tc["id"] for tc in out[0]["tool_calls"]] == ["t2"]
|
||||
assert out[1]["tool_call_id"] == "t2"
|
||||
|
||||
def test_orphan_tool_message_removed(self):
|
||||
"""既有行为不回归:孤立 tool 消息(无 assistant)移除。"""
|
||||
comp = _make_compressor()
|
||||
msgs = [
|
||||
{"role": "user", "content": "q"},
|
||||
{"role": "tool", "tool_call_id": "ghost", "content": "x"},
|
||||
]
|
||||
out = comp._validate_message_structure(msgs)
|
||||
assert [m["role"] for m in out] == ["user"]
|
||||
|
||||
def test_intact_pair_untouched(self):
|
||||
"""完整配对原样保留(同一对象引用,零开销路径)。"""
|
||||
comp = _make_compressor()
|
||||
assistant = {"role": "assistant", "content": "ok",
|
||||
"tool_calls": [{"id": "t1", "type": "function",
|
||||
"function": {"name": "a", "arguments": "{}"}}]}
|
||||
tool = {"role": "tool", "tool_call_id": "t1", "content": "r"}
|
||||
out = comp._validate_message_structure([assistant, tool])
|
||||
assert out == [assistant, tool]
|
||||
|
||||
|
||||
# ==================== T10:调度器硬超时兜底 ====================
|
||||
|
||||
class _UnkillableAI:
|
||||
"""软取消无法中断的卡死 AI(模拟 requests 读阻塞被滴流输出规避)。"""
|
||||
|
||||
def __init__(self):
|
||||
self.cancel_called = 0
|
||||
self._release = threading.Event()
|
||||
|
||||
def cancel(self):
|
||||
self.cancel_called += 1
|
||||
# 无效:think_and_respond 仍不返回
|
||||
|
||||
def think_and_respond(self, input_text, output_callback=None,
|
||||
is_continuation=False):
|
||||
self._release.wait(timeout=30)
|
||||
return "never-reached"
|
||||
|
||||
|
||||
class TestT10SchedulerHardTimeout:
|
||||
def test_stuck_ai_hard_timeout(self):
|
||||
"""软取消无效的卡死任务 → wait_for 硬超时放行,调度循环继续。"""
|
||||
from aibox import ConfigManager, TaskScheduler
|
||||
|
||||
class _FakeMemoDB:
|
||||
def get_memo(self, memo_id):
|
||||
return None
|
||||
|
||||
loop = asyncio.new_event_loop()
|
||||
fake_ai = _UnkillableAI()
|
||||
try:
|
||||
scheduler = TaskScheduler(loop=loop, logger=None)
|
||||
memo = {"id": 7, "title": "卡死任务", "content": "内容", "trigger_count": 0}
|
||||
orig_get = ConfigManager.get
|
||||
|
||||
def fake_get(self, key, default=None):
|
||||
# task_timeout 设为 -29.5:软取消立即触发(call_later 负延迟),
|
||||
# wait_for 硬超时 = -29.5 + 30 = 0.5s 后放行
|
||||
if key == "scheduler.task_timeout":
|
||||
return -29.5
|
||||
return orig_get(self, key, default)
|
||||
|
||||
with mock.patch.object(ConfigManager, "get", fake_get), \
|
||||
mock.patch.object(scheduler, "_create_task_ai", return_value=fake_ai):
|
||||
result = loop.run_until_complete(
|
||||
scheduler.execute_task_as_conversation(memo, _FakeMemoDB()))
|
||||
assert "硬超时" in result
|
||||
assert fake_ai.cancel_called >= 1 # 软取消先于硬超时触发
|
||||
finally:
|
||||
fake_ai._release.set() # 释放卡死线程
|
||||
loop.close()
|
||||
|
||||
def test_normal_task_unaffected(self):
|
||||
"""正常返回的任务不受硬超时包装影响(结果正常透传)。"""
|
||||
from aibox import ConfigManager, TaskScheduler
|
||||
|
||||
class _FakeMemoDB:
|
||||
def get_memo(self, memo_id):
|
||||
return None
|
||||
|
||||
class _FastAI:
|
||||
def cancel(self):
|
||||
pass
|
||||
|
||||
def think_and_respond(self, input_text, output_callback=None,
|
||||
is_continuation=False):
|
||||
return "任务完成"
|
||||
|
||||
loop = asyncio.new_event_loop()
|
||||
try:
|
||||
scheduler = TaskScheduler(loop=loop, logger=None)
|
||||
memo = {"id": 8, "title": "正常任务", "content": "内容", "trigger_count": 0}
|
||||
orig_get = ConfigManager.get
|
||||
|
||||
def fake_get(self, key, default=None):
|
||||
if key == "scheduler.task_timeout":
|
||||
return 600
|
||||
return orig_get(self, key, default)
|
||||
|
||||
with mock.patch.object(ConfigManager, "get", fake_get), \
|
||||
mock.patch.object(scheduler, "_create_task_ai", return_value=_FastAI()):
|
||||
result = loop.run_until_complete(
|
||||
scheduler.execute_task_as_conversation(memo, _FakeMemoDB()))
|
||||
assert "任务完成" in result
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
|
||||
# ==================== T22:owner 隔离(ProcessManager) ====================
|
||||
|
||||
class TestT22OwnerIsolation:
|
||||
def test_kill_denied_for_other_owner(self):
|
||||
from subsystems.process_manager import ProcessManager
|
||||
pm = ProcessManager()
|
||||
r = pm.start("sleep 30", owner="session-a")
|
||||
handle = r["handle"]
|
||||
try:
|
||||
denied = pm.kill(handle, owner="session-b")
|
||||
assert not denied["ok"]
|
||||
assert "属主" in denied["error"]
|
||||
# 进程仍在运行(未被误杀)
|
||||
assert pm._lookup(handle).status == "running"
|
||||
finally:
|
||||
pm.kill(handle)
|
||||
|
||||
def test_kill_allowed_for_owner(self):
|
||||
from subsystems.process_manager import ProcessManager
|
||||
pm = ProcessManager()
|
||||
r = pm.start("sleep 30", owner="session-a")
|
||||
result = pm.kill(r["handle"], owner="session-a")
|
||||
assert result["ok"]
|
||||
|
||||
def test_write_denied_for_other_owner(self):
|
||||
from subsystems.process_manager import ProcessManager
|
||||
pm = ProcessManager()
|
||||
r = pm.start("cat", owner="session-a")
|
||||
handle = r["handle"]
|
||||
try:
|
||||
denied = pm.write(handle, "hi", owner="session-b")
|
||||
assert not denied["ok"]
|
||||
assert "属主" in denied["error"]
|
||||
ok = pm.write(handle, "hi", owner="session-a")
|
||||
assert ok["ok"]
|
||||
finally:
|
||||
pm.kill(handle)
|
||||
|
||||
def test_backward_compatible_without_owner(self):
|
||||
"""任一侧无 owner 标识 → 不隔离(测试/管理面直调兼容)。"""
|
||||
from subsystems.process_manager import ProcessManager
|
||||
pm = ProcessManager()
|
||||
r1 = pm.start("sleep 5", owner="a")
|
||||
assert pm.kill(r1["handle"])["ok"] # caller 无 owner
|
||||
r2 = pm.start("sleep 5") # entry 无 owner
|
||||
assert pm.kill(r2["handle"], owner="b")["ok"]
|
||||
|
||||
def test_list_shows_owner_field(self):
|
||||
from subsystems.process_manager import ProcessManager
|
||||
pm = ProcessManager()
|
||||
r = pm.start("sleep 5", owner="whoami")
|
||||
try:
|
||||
infos = pm.list_procs()
|
||||
assert any(e.get("owner") == "whoami" for e in infos)
|
||||
finally:
|
||||
pm.kill(r["handle"])
|
||||
|
||||
def test_output_not_isolated(self):
|
||||
"""观察面(get_output)不受 owner 限制。"""
|
||||
from subsystems.process_manager import ProcessManager
|
||||
pm = ProcessManager()
|
||||
r = pm.start("echo hello", owner="session-a")
|
||||
handle = r["handle"]
|
||||
import time as _t
|
||||
_t.sleep(0.3)
|
||||
result = pm.get_output(handle) # 观察面无 owner 概念
|
||||
assert result["ok"]
|
||||
assert "hello" in result.get("output", "")
|
||||
|
||||
|
||||
# ==================== T22:task profile(定时任务精简工具面) ====================
|
||||
|
||||
def _make_chat(profile, owner=None):
|
||||
from aibox import DeepSeekChat
|
||||
with mock.patch("aibox.create_llm_client") as mclc:
|
||||
mclc.return_value = mock.MagicMock(
|
||||
api_key="sk-test", name="DS", default_temperature=None)
|
||||
return DeepSeekChat(tool_profile=profile, owner_id=owner)
|
||||
|
||||
|
||||
class TestT22TaskProfile:
|
||||
def test_task_profile_excludes_global_side_effect_tools(self):
|
||||
ai = _make_chat("task")
|
||||
names = set(ai.tool_registry._tools.keys())
|
||||
assert "update_identity" not in names
|
||||
assert "update_soul" not in names
|
||||
assert "set_proxy" not in names
|
||||
# 管理面其余工具保留
|
||||
assert "get_proxy" in names
|
||||
assert "create_subagent" in names
|
||||
|
||||
def test_orchestrator_profile_still_has_them(self):
|
||||
ai = _make_chat("orchestrator")
|
||||
names = set(ai.tool_registry._tools.keys())
|
||||
assert "update_identity" in names
|
||||
assert "update_soul" in names
|
||||
assert "set_proxy" in names
|
||||
|
||||
def test_worker_bg_tools_carry_owner(self):
|
||||
ai = _make_chat("worker", owner="sess-1")
|
||||
assert ai.tool_registry.get_tool("bg_run").owner == "sess-1"
|
||||
assert ai.tool_registry.get_tool("bg_kill").owner == "sess-1"
|
||||
assert ai.tool_registry.get_tool("bg_write").owner == "sess-1"
|
||||
# 观察面工具不带 owner
|
||||
assert getattr(ai.tool_registry.get_tool("bg_list"), "owner", "unset") == "unset"
|
||||
|
||||
def test_create_task_ai_uses_task_profile(self):
|
||||
from aibox import TaskScheduler
|
||||
# 显式新建 loop:全量运行时主线程可能已无可用的当前事件循环
|
||||
loop = asyncio.new_event_loop()
|
||||
try:
|
||||
scheduler = TaskScheduler(loop=loop, logger=None)
|
||||
with mock.patch("aibox.create_llm_client") as mclc:
|
||||
mclc.return_value = mock.MagicMock(
|
||||
api_key="sk-test", name="DS", default_temperature=None)
|
||||
ai = scheduler._create_task_ai()
|
||||
assert ai.tool_profile == "task"
|
||||
assert ai.owner_id and ai.owner_id.startswith("task-")
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
|
||||
# ==================== T-L17/T-L18:model_config 展示与推荐 ====================
|
||||
|
||||
class TestTL17EnvSuggestion:
|
||||
def test_suggest_by_domain(self):
|
||||
from model_config import _suggest_env_name
|
||||
assert _suggest_env_name("https://api.moonshot.cn/v1") == "MOONSHOT_API_KEY"
|
||||
assert _suggest_env_name("https://api.deepseek.com/chat") == "DEEPSEEK_API_KEY"
|
||||
assert _suggest_env_name("https://dashscope.aliyuncs.com/x") == "DASHSCOPE_API_KEY"
|
||||
|
||||
def test_unknown_domain_empty(self):
|
||||
from model_config import _suggest_env_name
|
||||
assert _suggest_env_name("http://localhost:11434/v1") == ""
|
||||
assert _suggest_env_name("") == ""
|
||||
assert _suggest_env_name(None) == ""
|
||||
|
||||
|
||||
class TestTL18DisplaySanitize:
|
||||
def test_sanitize_strips_control_chars(self):
|
||||
from model_config import _sanitize_display
|
||||
out = _sanitize_display("a\x1b]0;pwned\x07b")
|
||||
assert "\x1b" not in out and "\x07" not in out
|
||||
assert out == "a]0;pwnedb" # 控制符剥离,可打印部分保留
|
||||
assert _sanitize_display("正常文本") == "正常文本"
|
||||
assert _sanitize_display(None) == ""
|
||||
|
||||
def test_print_profiles_filters_escape_sequences(self, capsys):
|
||||
from model_config import _print_profiles
|
||||
_print_profiles([{
|
||||
"name": "evil\x1b[31m", "is_default": False,
|
||||
"display_name": "d\x07x", "api_url": "https://x\x1b[2J",
|
||||
"model": "m", "api_key_env": "K\x1bZ",
|
||||
"temperature": None, "timeout": None, "fallback": ["fb\x07"],
|
||||
}])
|
||||
out = capsys.readouterr().out
|
||||
assert "\x1b" not in out
|
||||
assert "\x07" not in out
|
||||
assert "evil" in out # 可打印部分保留
|
||||
@@ -0,0 +1,448 @@
|
||||
"""后台进程管理器(subsystems.process_manager)与 bg_* 工具回归测试
|
||||
|
||||
覆盖三决策语义与安全边界:
|
||||
- bg 生命周期:启动立即返回 / 惰性收割 / kill 幂等 / TTL 过期
|
||||
- 输出缓冲:tail/since 游标续读 / 环形淘汰提示 / 单行截断 / 拉取字节上限
|
||||
- 上限:运行中进程数上限 / 总条目 LRU 淘汰
|
||||
- 收尾:shutdown 杀非 detach、保留 detach;atexit 单例共享
|
||||
- 工具层:参数规整(字符串句柄)/ 双闸禁用 / 注册表分发
|
||||
|
||||
跨平台:长驻进程用 sys.executable -c "time.sleep(...)"(echo/ping 不可移植)。
|
||||
Windows 注意:shell=True 下 terminate 只杀 shell 包装进程(与执行器 L9
|
||||
回退语义一致),被杀判定以 popen.poll() 为准。
|
||||
"""
|
||||
|
||||
import queue
|
||||
import sys
|
||||
import time
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
if "websockets" not in sys.modules:
|
||||
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||
|
||||
from subsystems import ProcessManager, get_process_manager, reset_process_manager # noqa: E402
|
||||
from subsystems.process_manager import MAX_LINE_BYTES # noqa: E402
|
||||
from tools import ( # noqa: E402
|
||||
ToolRegistry,
|
||||
BgRunTool, BgWriteTool, BgListTool, BgOutputTool, BgKillTool,
|
||||
)
|
||||
|
||||
|
||||
PY = sys.executable
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def echo_srv(tmp_path):
|
||||
"""行协议回环服务器:stdin 读一行 → stdout 回 ECHO:行,EOF 退出。
|
||||
交互式 I/O(bg_write/watermark/优雅终止)测试的标准对端。"""
|
||||
script = (
|
||||
"import sys\n"
|
||||
"while True:\n"
|
||||
" line = sys.stdin.readline()\n"
|
||||
" if not line:\n"
|
||||
" break\n"
|
||||
" sys.stdout.write('ECHO:' + line.strip() + '\\n')\n"
|
||||
" sys.stdout.flush()\n"
|
||||
)
|
||||
f = tmp_path / "echo_srv.py"
|
||||
f.write_text(script, encoding="utf-8")
|
||||
return f'"{PY}" "{f}"'
|
||||
|
||||
|
||||
def _py(code: str) -> str:
|
||||
"""跨平台 python 命令(code 内只用单引号)"""
|
||||
return f'"{PY}" -c "{code}"'
|
||||
|
||||
|
||||
ECHO_10 = _py("for i in range(10): print('line-%d' % i)")
|
||||
SLEEP_10 = _py("import time; time.sleep(10)")
|
||||
HELLO = _py("print('hello-bg')")
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mgr():
|
||||
"""每测试独立管理器(不走单例),收尾兜底清进程"""
|
||||
m = ProcessManager()
|
||||
try:
|
||||
yield m
|
||||
finally:
|
||||
m.shutdown()
|
||||
|
||||
|
||||
def _wait_status(m: ProcessManager, handle: int, status: str, timeout: float = 10.0):
|
||||
"""轮询等待句柄达到指定状态"""
|
||||
deadline = time.time() + timeout
|
||||
while time.time() < deadline:
|
||||
for p in m.list_procs():
|
||||
if p["handle"] == handle and p["status"] == status:
|
||||
return p
|
||||
time.sleep(0.1)
|
||||
return None
|
||||
|
||||
|
||||
# ==================== 启动与生命周期 ====================
|
||||
|
||||
class TestStart:
|
||||
def test_start_returns_handle_immediately(self, mgr):
|
||||
t0 = time.time()
|
||||
r = mgr.start(SLEEP_10, name="sleeper")
|
||||
assert r["ok"] is True
|
||||
assert time.time() - t0 < 2 # 立即返回,不等进程结束
|
||||
assert isinstance(r["handle"], int) and r["pid"] > 0
|
||||
assert r["name"] == "sleeper"
|
||||
|
||||
def test_start_default_name_from_command(self, mgr):
|
||||
r = mgr.start(HELLO)
|
||||
assert r["ok"] is True
|
||||
# 默认名称 = 命令前 40 字符
|
||||
assert r["name"] == HELLO.strip()[:40]
|
||||
|
||||
def test_start_rejects_empty_command(self, mgr):
|
||||
assert mgr.start(" ")["ok"] is False
|
||||
assert mgr.start(None)["ok"] is False
|
||||
|
||||
def test_start_rejects_bad_name(self, mgr):
|
||||
# 名称含控制字符(M5 惯例:防终端/前端注入)
|
||||
assert mgr.start(HELLO, name="bad\x1b[31m")["ok"] is False
|
||||
assert mgr.start(HELLO, name="x" * 65)["ok"] is False
|
||||
|
||||
def test_start_missing_binary_returns_error(self, mgr):
|
||||
r = mgr.start("definitely-not-a-command-xyz --version")
|
||||
# shell 本身能启动(rc=127/9009),故启动成功、退出码非零
|
||||
assert r["ok"] is True
|
||||
p = _wait_status(mgr, r["handle"], "exited")
|
||||
assert p is not None
|
||||
assert p["exit_code"] != 0
|
||||
|
||||
|
||||
class TestLifecycle:
|
||||
def test_exited_status_and_exit_code(self, mgr):
|
||||
r = mgr.start(HELLO)
|
||||
p = _wait_status(mgr, r["handle"], "exited")
|
||||
assert p is not None
|
||||
assert p["exit_code"] == 0
|
||||
|
||||
def test_kill_running_process(self, mgr):
|
||||
r = mgr.start(SLEEP_10)
|
||||
time.sleep(0.5) # 等进程真正起来
|
||||
k = mgr.kill(r["handle"])
|
||||
assert k["ok"] is True and k["status"] == "killed"
|
||||
assert mgr._procs[r["handle"]].popen.poll() is not None
|
||||
|
||||
def test_kill_is_idempotent_on_exited(self, mgr):
|
||||
r = mgr.start(HELLO)
|
||||
assert _wait_status(mgr, r["handle"], "exited") is not None
|
||||
k = mgr.kill(r["handle"])
|
||||
assert k["ok"] is True and k["already_finished"] is True
|
||||
|
||||
def test_kill_missing_handle(self, mgr):
|
||||
assert mgr.kill(999)["ok"] is False
|
||||
|
||||
def test_ttl_expiry(self):
|
||||
m = ProcessManager(max_lifetime=0.5)
|
||||
try:
|
||||
r = m.start(SLEEP_10)
|
||||
time.sleep(0.8)
|
||||
# list 触发惰性 sweep
|
||||
assert _wait_status(m, r["handle"], "expired", timeout=2) is not None
|
||||
finally:
|
||||
m.shutdown()
|
||||
|
||||
def test_max_processes_limit(self):
|
||||
m = ProcessManager(max_processes=2)
|
||||
try:
|
||||
assert m.start(SLEEP_10)["ok"] is True
|
||||
assert m.start(SLEEP_10)["ok"] is True
|
||||
r3 = m.start(SLEEP_10)
|
||||
assert r3["ok"] is False and "上限" in r3["error"]
|
||||
finally:
|
||||
m.shutdown()
|
||||
|
||||
def test_max_entries_evicts_oldest_exited(self):
|
||||
m = ProcessManager(max_entries=2)
|
||||
try:
|
||||
h1 = m.start(HELLO)["handle"]
|
||||
h2 = m.start(HELLO)["handle"]
|
||||
assert _wait_status(m, h1, "exited") is not None
|
||||
assert _wait_status(m, h2, "exited") is not None
|
||||
# 第三条触发最老已退出条目(#1)LRU 淘汰
|
||||
h3 = m.start(HELLO)["handle"]
|
||||
assert h3 > 0
|
||||
handles = {p["handle"] for p in m.list_procs()}
|
||||
assert h1 not in handles and h2 in handles and h3 in handles
|
||||
finally:
|
||||
m.shutdown()
|
||||
|
||||
|
||||
# ==================== 输出缓冲与游标 ====================
|
||||
|
||||
class TestOutput:
|
||||
def test_tail_and_since_cursor(self, mgr):
|
||||
r = mgr.start(ECHO_10)
|
||||
assert _wait_status(mgr, r["handle"], "exited") is not None
|
||||
time.sleep(0.3) # reader 排空末尾行
|
||||
|
||||
out = mgr.get_output(r["handle"], tail=3)
|
||||
assert out["ok"] is True
|
||||
assert out["lines"] == 3
|
||||
assert "line-9" in out["output"] and "line-6" not in out["output"]
|
||||
|
||||
full = mgr.get_output(r["handle"], since=0)
|
||||
assert full["lines"] == 10
|
||||
assert full["next_cursor"] == 10
|
||||
|
||||
inc = mgr.get_output(r["handle"], since=full["next_cursor"])
|
||||
assert inc["lines"] == 0 and inc["next_cursor"] == 10
|
||||
|
||||
def test_since_before_eviction_reports_gap(self):
|
||||
m = ProcessManager(buffer_per_process=2048)
|
||||
try:
|
||||
r = m.start(_py("for i in range(50): print('pad-%03d-%s' % (i, 'x'*80))"))
|
||||
assert _wait_status(m, r["handle"], "exited") is not None
|
||||
time.sleep(0.3)
|
||||
out = m.get_output(r["handle"], since=0)
|
||||
assert out["ok"] is True
|
||||
assert out["evicted"] > 0 # 早期行被环形淘汰
|
||||
assert out["lines"] < 50
|
||||
finally:
|
||||
m.shutdown()
|
||||
|
||||
def test_single_line_truncated(self, mgr):
|
||||
r = mgr.start(_py("print('A' * (1024*64))"))
|
||||
assert _wait_status(mgr, r["handle"], "exited") is not None
|
||||
time.sleep(0.3)
|
||||
out = mgr.get_output(r["handle"], tail=5)
|
||||
assert "单行截断" in out["output"]
|
||||
# 缓冲内存预算不被超长单行击穿
|
||||
entry = mgr._procs[r["handle"]]
|
||||
assert entry.buffer_bytes <= MAX_LINE_BYTES + 1024
|
||||
|
||||
def test_fetch_byte_limit_truncates(self, mgr):
|
||||
r = mgr.start(_py("for i in range(100): print('y' * 200)"))
|
||||
assert _wait_status(mgr, r["handle"], "exited") is not None
|
||||
time.sleep(0.3)
|
||||
out = mgr.get_output(r["handle"], since=0, max_bytes=4096)
|
||||
assert out["truncated"] is True
|
||||
|
||||
def test_missing_handle(self, mgr):
|
||||
assert mgr.get_output(999)["ok"] is False
|
||||
|
||||
def test_invalid_since(self, mgr):
|
||||
r = mgr.start(HELLO)
|
||||
assert mgr.get_output(r["handle"], since="abc")["ok"] is False
|
||||
|
||||
|
||||
# ==================== 交互式 I/O(stdin 管道 + watermark 对齐) ====================
|
||||
|
||||
class TestInteractive:
|
||||
def test_write_roundtrip_watermark_alignment(self, mgr, echo_srv):
|
||||
"""核心对齐语义:bg_output(since=watermark) = 本条输入的响应且不含回显"""
|
||||
h = mgr.start(echo_srv)["handle"]
|
||||
r = mgr.write(h, "hello", wait_lines=1, wait_timeout=10)
|
||||
assert r["ok"] is True
|
||||
assert "ECHO:hello" in r["wait"]["output"]
|
||||
assert r["wait"]["timed_out"] is False
|
||||
|
||||
out = mgr.get_output(h, since=r["watermark"])
|
||||
assert "ECHO:hello" in out["output"]
|
||||
assert "[stdin]" not in out["output"] # 回显行序号 == watermark,不包含
|
||||
|
||||
def test_echo_line_visible_in_tail(self, mgr, echo_srv):
|
||||
"""[stdin] 回显进环形缓冲,tail 视图可见完整交互时序"""
|
||||
h = mgr.start(echo_srv)["handle"]
|
||||
assert mgr.write(h, "abc", wait_lines=1)["ok"] is True
|
||||
out = mgr.get_output(h, tail=10)
|
||||
assert "[stdin] abc" in out["output"]
|
||||
assert "ECHO:abc" in out["output"]
|
||||
|
||||
def test_write_wait_sync_roundtrip(self, mgr, echo_srv):
|
||||
"""wait_lines 同步等待:一步问答(expect 式)"""
|
||||
h = mgr.start(echo_srv)["handle"]
|
||||
r = mgr.write(h, "1+1", wait_lines=1)
|
||||
assert r["wait"]["lines"] == 1
|
||||
assert r["wait"]["output"].strip() == "ECHO:1+1"
|
||||
assert r["wait"]["next_cursor"] == r["watermark"] + 1
|
||||
|
||||
def test_write_wait_timeout_on_silent_process(self, mgr):
|
||||
"""不读不响应的进程:写入成功(管道缓冲),等待如实超时"""
|
||||
h = mgr.start(SLEEP_10)["handle"]
|
||||
r = mgr.write(h, "ping", wait_lines=1, wait_timeout=0.5)
|
||||
assert r["ok"] is True
|
||||
assert r["wait"]["timed_out"] is True
|
||||
assert r["wait"]["lines"] == 0
|
||||
|
||||
def test_write_to_exited_process(self, mgr):
|
||||
h = mgr.start(HELLO)["handle"]
|
||||
assert _wait_status(mgr, h, "exited") is not None
|
||||
assert mgr.write(h, "x")["ok"] is False
|
||||
|
||||
def test_write_rejects_empty_and_oversize(self, mgr, echo_srv):
|
||||
h = mgr.start(echo_srv)["handle"]
|
||||
assert mgr.write(h, "")["ok"] is False
|
||||
assert mgr.write(h, None)["ok"] is False
|
||||
assert mgr.write(h, "A" * (64 * 1024 + 1))["ok"] is False
|
||||
|
||||
def test_stdin_eof_graceful_kill(self, mgr, echo_srv):
|
||||
"""三级优雅终止第一级:关 stdin → 回环服务器见 EOF 自行退出"""
|
||||
h = mgr.start(echo_srv)["handle"]
|
||||
time.sleep(0.5)
|
||||
k = mgr.kill(h) # 不 force:先关 stdin
|
||||
assert k["ok"] is True
|
||||
assert mgr._procs[h].popen.poll() is not None
|
||||
assert _wait_status(mgr, h, "killed") is not None
|
||||
|
||||
def test_write_after_kill_rejected(self, mgr, echo_srv):
|
||||
h = mgr.start(echo_srv)["handle"]
|
||||
assert mgr.kill(h)["ok"] is True
|
||||
assert mgr.write(h, "late")["ok"] is False
|
||||
|
||||
def test_log_file_tee(self, mgr, monkeypatch, tmp_path):
|
||||
"""log_file 全档镜像:输出与 [stdin] 回显均落盘,可被原生工具读取"""
|
||||
from common import constants as cm
|
||||
monkeypatch.setattr(type(cm), "WORKSPACE_DIR",
|
||||
property(lambda self: str(tmp_path)))
|
||||
script = tmp_path / "srv.py"
|
||||
script.write_text(
|
||||
"import sys\n"
|
||||
"for line in sys.stdin:\n"
|
||||
" sys.stdout.write('ECHO:' + line.strip() + '\\n')\n"
|
||||
" sys.stdout.flush()\n",
|
||||
encoding="utf-8")
|
||||
cmd = f'"{PY}" "{script}"'
|
||||
r = mgr.start(cmd, log_file=True)
|
||||
assert r["ok"] is True and r["log_file"] is not None
|
||||
assert mgr.write(r["handle"], "tee-me", wait_lines=1)["ok"] is True
|
||||
mgr.kill(r["handle"])
|
||||
log = (tmp_path / f"bg_{r['handle']}.log").read_text(encoding="utf-8")
|
||||
assert "[stdin] tee-me" in log
|
||||
assert "ECHO:tee-me" in log
|
||||
|
||||
def test_bg_write_tool_flow(self, mgr, echo_srv):
|
||||
"""工具层冒烟:BgWriteTool 参数规整 + 对齐读取提示"""
|
||||
assert BgRunTool(manager=mgr).execute(command=echo_srv).startswith("✅")
|
||||
out = BgWriteTool(manager=mgr).execute(handle="1", text="hi",
|
||||
wait_lines="1", wait_timeout=10)
|
||||
assert "ECHO:hi" in out
|
||||
assert "watermark" in out
|
||||
# 描述里的行缓冲指引(给模型看的适用前提)
|
||||
assert "PYTHONUNBUFFERED" in BgWriteTool(manager=mgr).get_description()
|
||||
|
||||
def test_registry_dispatch_bg_write(self, mgr, echo_srv):
|
||||
reg = ToolRegistry()
|
||||
reg.register_many([BgRunTool(manager=mgr), BgWriteTool(manager=mgr)],
|
||||
is_builtin=True)
|
||||
assert reg.execute_tool("bg_run", command=echo_srv).startswith("✅")
|
||||
r = reg.execute_tool("bg_write", handle=1, text="via-reg",
|
||||
wait_lines=1, wait_timeout=10)
|
||||
assert "ECHO:via-reg" in r
|
||||
|
||||
|
||||
# ==================== 退出收尾 ====================
|
||||
|
||||
class TestShutdown:
|
||||
def test_shutdown_kills_non_detach_keeps_detach(self):
|
||||
m = ProcessManager()
|
||||
try:
|
||||
h1 = m.start(SLEEP_10)["handle"] # 常规:退出时终止
|
||||
h2 = m.start(SLEEP_10, detach=True)["handle"] # detach:保留
|
||||
time.sleep(0.5)
|
||||
e1, e2 = m._procs[h1], m._procs[h2]
|
||||
|
||||
r = m.shutdown()
|
||||
assert r["killed"] == 1 and r["detached"] == 1
|
||||
assert e1.popen.poll() is not None # 常规进程已终止
|
||||
assert e2.popen.poll() is None # detach 进程仍存活
|
||||
assert m.list_procs() == [] # 注册表已清空(detach 仅注销)
|
||||
|
||||
# 幂等
|
||||
assert m.shutdown() == {"killed": 0, "detached": 0}
|
||||
# 收尾后拒绝新进程
|
||||
assert m.start(HELLO)["ok"] is False
|
||||
# 清理 detach 测试进程
|
||||
e2.popen.kill()
|
||||
finally:
|
||||
m.shutdown()
|
||||
|
||||
def test_singleton_shared_across_callers(self):
|
||||
m1 = get_process_manager()
|
||||
m2 = get_process_manager()
|
||||
assert m1 is m2 # 跨会话/跨子代理共享同一注册表
|
||||
reset_process_manager(None)
|
||||
m3 = get_process_manager()
|
||||
assert m3 is not m1
|
||||
m3.shutdown()
|
||||
reset_process_manager(None)
|
||||
|
||||
|
||||
# ==================== 工具层 ====================
|
||||
|
||||
class TestTools:
|
||||
def test_bg_run_and_output_flow(self, mgr):
|
||||
run = BgRunTool(manager=mgr)
|
||||
r = run.execute(command=HELLO, proc_name="hello")
|
||||
assert r.startswith("✅") and "#" in r
|
||||
handle = 1
|
||||
assert _wait_status(mgr, handle, "exited") is not None
|
||||
time.sleep(0.3)
|
||||
|
||||
out = BgOutputTool(manager=mgr).execute(handle=handle, since=0)
|
||||
assert "hello-bg" in out
|
||||
assert "next_cursor" in out
|
||||
|
||||
def test_bg_run_string_handle_coercion(self, mgr):
|
||||
"""LLM 可能传字符串数字句柄(W 批输入规整惯例)"""
|
||||
run = BgRunTool(manager=mgr)
|
||||
assert run.execute(command=HELLO).startswith("✅")
|
||||
assert _wait_status(mgr, 1, "exited") is not None
|
||||
out = BgOutputTool(manager=mgr).execute(handle="1", tail=5)
|
||||
assert "hello-bg" in out
|
||||
|
||||
def test_bg_run_missing_command(self, mgr):
|
||||
r = BgRunTool(manager=mgr).execute()
|
||||
assert "缺少" in r
|
||||
|
||||
def test_bg_list_empty_and_filled(self, mgr):
|
||||
lst = BgListTool(manager=mgr)
|
||||
assert "没有后台进程" in lst.execute()
|
||||
mgr.start(SLEEP_10, name="sleeper")
|
||||
r = lst.execute()
|
||||
assert "sleeper" in r and "运行中 1" in r
|
||||
|
||||
def test_bg_kill_flow(self, mgr):
|
||||
mgr.start(SLEEP_10, name="sleeper")
|
||||
time.sleep(0.5)
|
||||
r = BgKillTool(manager=mgr).execute(handle=1)
|
||||
assert r.startswith("✅") and "已终止" in r
|
||||
|
||||
def test_disabled_gate_blocks_all(self, mgr, monkeypatch):
|
||||
from common import constants
|
||||
monkeypatch.setattr(type(constants), "BG_PROCESSES_ENABLED",
|
||||
property(lambda self: False))
|
||||
assert "禁用" in BgRunTool(manager=mgr).execute(command=HELLO)
|
||||
assert "禁用" in BgListTool(manager=mgr).execute()
|
||||
assert "禁用" in BgOutputTool(manager=mgr).execute(handle=1)
|
||||
assert "禁用" in BgKillTool(manager=mgr).execute(handle=1)
|
||||
|
||||
def test_registry_dispatch(self, mgr):
|
||||
"""4 个工具经 ToolRegistry 正常分发(worker 注册链路冒烟)"""
|
||||
reg = ToolRegistry()
|
||||
reg.register_many([
|
||||
BgRunTool(manager=mgr), BgListTool(manager=mgr),
|
||||
BgOutputTool(manager=mgr), BgKillTool(manager=mgr),
|
||||
], is_builtin=True)
|
||||
r = reg.execute_tool("bg_run", command=HELLO, proc_name="smoke")
|
||||
assert r.startswith("✅")
|
||||
assert _wait_status(mgr, 1, "exited") is not None
|
||||
time.sleep(0.3)
|
||||
r = reg.execute_tool("bg_output", handle=1, since=0)
|
||||
assert "hello-bg" in r
|
||||
r = reg.execute_tool("bg_list")
|
||||
assert "smoke" in r
|
||||
|
||||
|
||||
def test_worker_only_tools_includes_bg():
|
||||
"""编排者误调拦截表包含 bg 工具(worker 执行面边界)"""
|
||||
from aibox import WORKER_ONLY_TOOLS
|
||||
assert {"bg_run", "bg_write", "bg_list", "bg_output", "bg_kill"} <= WORKER_ONLY_TOOLS
|
||||
@@ -0,0 +1,178 @@
|
||||
"""run_full_subagent 最大轮次(max_turns)行为回归测试(2026-08-19)
|
||||
|
||||
覆盖 aibox.py run_full_subagent 的轮次限制链路(think_with_limit 补丁):
|
||||
- 超限拦截:任务轮 LLM 调用被限制在 max_turns 次,超限递归返回哨兵串
|
||||
- 哨兵串文本契约("已达到最大对话轮次,任务结束。")
|
||||
- 收尾报告:超限/正常结束均由 original_think 绕过轮次限制生成,作为返回值
|
||||
- 入参钳制:max_turns 为 0/None 回退默认 10,负数钳制下限 1
|
||||
- 超限不置 conversation_active=False(收尾保留完整上下文,aibox.py 763-765 注释)
|
||||
- 收尾失败兜底链:最后一条 assistant 消息 → "子代理未产生有效对话内容。"
|
||||
- SubAgentTool.execute 的 max_turns 参数透传(默认 10)
|
||||
"""
|
||||
|
||||
import sys
|
||||
import threading
|
||||
import types
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
# ---- 环境桩:websockets 未安装时保证 import aibox.py 不失败 ----
|
||||
if "websockets" not in sys.modules:
|
||||
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||
|
||||
from aibox import SubAgentTool, run_full_subagent # noqa: E402
|
||||
|
||||
# 与 aibox.py think_with_limit 超限返回值一致的哨兵串(文本契约)
|
||||
SENTINEL = "已达到最大对话轮次,任务结束。"
|
||||
WRAP_UP_MARKER = "【收尾】"
|
||||
|
||||
|
||||
def _make_fake_chat(task_rounds=999, final_answer=None, wrap_report="收尾报告"):
|
||||
"""模拟 DeepSeekChat:think_and_respond 的"原始实现"。
|
||||
|
||||
每次任务轮调用递归一次(模拟 aibox.py:1801 工具调用后递归继续,
|
||||
经由补丁后的 think_and_respond 走轮次计数);累计 task_rounds 轮后
|
||||
若给出 final_answer 则正常结束(模拟 AI 返回纯文本)。task_rounds
|
||||
设为大数且不给 final_answer 时永不自行收敛——只有轮次限制能终止它。
|
||||
检测到收尾 prompt(含【收尾】标记)时返回 wrap_report。
|
||||
"""
|
||||
fake = mock.MagicMock()
|
||||
state = {"n": 0, "task_inputs": [], "recurse_returns": []}
|
||||
|
||||
def original_think(input_text, output_callback=None, is_continuation=False):
|
||||
if WRAP_UP_MARKER in input_text:
|
||||
return wrap_report
|
||||
state["task_inputs"].append(input_text)
|
||||
state["n"] += 1
|
||||
if final_answer is not None and state["n"] >= task_rounds:
|
||||
return final_answer
|
||||
ret = fake.think_and_respond("请基于以上工具执行结果继续回答。",
|
||||
output_callback, is_continuation=True)
|
||||
state["recurse_returns"].append(ret)
|
||||
return ret
|
||||
|
||||
fake.think_and_respond = original_think
|
||||
fake.stats = {}
|
||||
fake.messages = []
|
||||
fake._cancel_event = threading.Event()
|
||||
fake.conversation_active = True
|
||||
fake._state = state
|
||||
return fake
|
||||
|
||||
|
||||
def _run(fake, **kwargs):
|
||||
with mock.patch("aibox.DeepSeekChat", return_value=fake), \
|
||||
mock.patch("aibox.create_llm_client"), \
|
||||
mock.patch("aibox.SubprocessCommandExecutor"):
|
||||
return run_full_subagent("角色", "任务", None, **kwargs)
|
||||
|
||||
|
||||
class TestMaxTurnsCutoff:
|
||||
def test_task_turns_capped_at_max_turns(self):
|
||||
"""永不收敛的任务:任务轮 LLM 调用恰为 max_turns 次,第 max_turns+1 次被拦截"""
|
||||
fake = _make_fake_chat(task_rounds=999)
|
||||
result = _run(fake, max_turns=3)
|
||||
assert len(fake._state["task_inputs"]) == 3
|
||||
assert result == "收尾报告"
|
||||
|
||||
def test_overrun_recursion_returns_sentinel(self):
|
||||
"""超限那轮递归返回哨兵串(不调 LLM),并沿递归链上抛"""
|
||||
fake = _make_fake_chat(task_rounds=999)
|
||||
_run(fake, max_turns=3)
|
||||
assert fake._state["recurse_returns"]
|
||||
assert all(r == SENTINEL for r in fake._state["recurse_returns"])
|
||||
|
||||
def test_wrap_up_runs_after_overrun(self):
|
||||
"""超限后收尾报告仍被请求(original_think 绕过轮次限制),且为最终返回值"""
|
||||
fake = _make_fake_chat(task_rounds=999, wrap_report="最终报告")
|
||||
result = _run(fake, max_turns=2)
|
||||
assert result == "最终报告"
|
||||
|
||||
def test_max_turns_zero_clamped_to_default(self):
|
||||
"""max_turns=0 → 回退默认 10(对齐 max_tokens 的 0/空回退惯例)"""
|
||||
fake = _make_fake_chat(task_rounds=999)
|
||||
_run(fake, max_turns=0)
|
||||
assert len(fake._state["task_inputs"]) == 10
|
||||
|
||||
def test_negative_max_turns_clamped_to_one(self):
|
||||
"""max_turns 为负 → 钳制下限 1(至少执行一个任务轮,不再首轮即拦截)"""
|
||||
fake = _make_fake_chat(task_rounds=999)
|
||||
_run(fake, max_turns=-3)
|
||||
assert len(fake._state["task_inputs"]) == 1
|
||||
|
||||
def test_max_turns_none_falls_back_to_default(self):
|
||||
"""max_turns=None → 默认 10"""
|
||||
fake = _make_fake_chat(task_rounds=999)
|
||||
_run(fake, max_turns=None)
|
||||
assert len(fake._state["task_inputs"]) == 10
|
||||
|
||||
def test_overrun_does_not_deactivate_conversation(self):
|
||||
"""超限不置 conversation_active=False(否则收尾前上下文被重置清空)"""
|
||||
fake = _make_fake_chat(task_rounds=999)
|
||||
_run(fake, max_turns=2)
|
||||
assert fake.conversation_active is True
|
||||
|
||||
|
||||
class TestNormalCompletion:
|
||||
def test_result_is_wrap_up_report_not_final_answer(self):
|
||||
"""正常收敛时返回值同样是收尾报告(最终回答不作为工具结果返回)"""
|
||||
fake = _make_fake_chat(task_rounds=2, final_answer="最终答案")
|
||||
result = _run(fake, max_turns=10)
|
||||
assert len(fake._state["task_inputs"]) == 2
|
||||
# 第一轮递归的返回值即第 2 轮的最终回答(无哨兵串、无超限拦截)
|
||||
assert fake._state["recurse_returns"] == ["最终答案"]
|
||||
assert result == "收尾报告"
|
||||
|
||||
def test_sentinel_never_appears_when_within_limit(self):
|
||||
"""未达上限:任何递归都不返回哨兵串"""
|
||||
fake = _make_fake_chat(task_rounds=3, final_answer="最终答案")
|
||||
_run(fake, max_turns=10)
|
||||
assert SENTINEL not in fake._state["recurse_returns"]
|
||||
|
||||
|
||||
class TestWrapUpFallback:
|
||||
def test_empty_wrap_up_falls_back_to_last_assistant_message(self):
|
||||
"""收尾返回空 → 兜底取最后一条有内容的 assistant 消息"""
|
||||
fake = _make_fake_chat(task_rounds=1, final_answer="x", wrap_report="")
|
||||
fake.messages = [
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "中间回答"},
|
||||
{"role": "assistant", "content": ""},
|
||||
{"role": "assistant", "content": "最终回答"},
|
||||
]
|
||||
assert _run(fake, max_turns=10) == "【子代理最终回答】\n最终回答"
|
||||
|
||||
def test_no_wrap_up_and_no_messages(self):
|
||||
"""收尾为空且无任何 assistant 消息 → 固定兜底文案"""
|
||||
fake = _make_fake_chat(task_rounds=1, final_answer="x", wrap_report="")
|
||||
assert _run(fake, max_turns=10) == "子代理未产生有效对话内容。"
|
||||
|
||||
def test_wrap_up_exception_falls_back(self):
|
||||
"""收尾调用抛异常 → 不中断,走 assistant 消息兜底"""
|
||||
fake = _make_fake_chat(task_rounds=1, final_answer="x")
|
||||
original = fake.think_and_respond
|
||||
|
||||
def boom_on_wrap_up(input_text, output_callback=None, is_continuation=False):
|
||||
if WRAP_UP_MARKER in input_text:
|
||||
raise RuntimeError("wrap-up boom")
|
||||
return original(input_text, output_callback, is_continuation)
|
||||
|
||||
fake.think_and_respond = boom_on_wrap_up
|
||||
fake.messages = [{"role": "assistant", "content": "最终回答"}]
|
||||
assert _run(fake, max_turns=10) == "【子代理最终回答】\n最终回答"
|
||||
|
||||
|
||||
class TestSubAgentToolMaxTurns:
|
||||
def test_max_turns_passed_through(self):
|
||||
"""SubAgentTool.execute 的 max_turns 透传给 run_full_subagent"""
|
||||
tool = SubAgentTool(logger=None)
|
||||
with mock.patch("aibox.run_full_subagent", return_value="ok") as m:
|
||||
tool.execute(role="r", task="t", max_turns=4)
|
||||
assert m.call_args.kwargs["max_turns"] == 4
|
||||
|
||||
def test_max_turns_default_is_10(self):
|
||||
tool = SubAgentTool(logger=None)
|
||||
with mock.patch("aibox.run_full_subagent", return_value="ok") as m:
|
||||
tool.execute(role="r", task="t")
|
||||
assert m.call_args.kwargs["max_turns"] == 10
|
||||
@@ -45,6 +45,13 @@ from .prompt_tools import (
|
||||
)
|
||||
from .proxy_tools import SetProxyTool, GetProxyTool
|
||||
from .misc_tools import AskUserTool
|
||||
from .process_tools import (
|
||||
BgRunTool,
|
||||
BgWriteTool,
|
||||
BgListTool,
|
||||
BgOutputTool,
|
||||
BgKillTool,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"Tool",
|
||||
@@ -71,4 +78,10 @@ __all__ = [
|
||||
"SetProxyTool",
|
||||
"GetProxyTool",
|
||||
"AskUserTool",
|
||||
# 后台进程管理工具(bg_run/bg_write/bg_list/bg_output/bg_kill)
|
||||
"BgRunTool",
|
||||
"BgWriteTool",
|
||||
"BgListTool",
|
||||
"BgOutputTool",
|
||||
"BgKillTool",
|
||||
]
|
||||
+196
-8
@@ -15,8 +15,10 @@ tools.basic_tools - 基础工具集(A 批抽离)
|
||||
依赖方向:common <- tools.basic_tools,不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
import base64
|
||||
import os
|
||||
import queue
|
||||
import re
|
||||
import shlex
|
||||
from datetime import datetime
|
||||
from typing import Dict, Optional, Tuple
|
||||
@@ -25,6 +27,40 @@ from common import ConfigManager, constants
|
||||
from .base import Tool
|
||||
|
||||
|
||||
# T3(审计 2026-08-21):cmd.exe 变量展开形态——%VAR% 匹配已定义环境变量时
|
||||
# 展开为变量值,值中的 & | 会成为命令分隔符(经典 cmd 注入)。URL 编码
|
||||
# (%20/%E4 等)与 PowerShell 的 %{http_code} 均不匹配该模式,不受影响。
|
||||
_WIN_VAR_PATTERN = re.compile(r"%[A-Za-z_][A-Za-z0-9_]*%")
|
||||
|
||||
|
||||
def _quote_arg(value: str) -> Tuple[str, Optional[str]]:
|
||||
"""平台感知的 shell 参数引用,返回 (quoted, error)。
|
||||
|
||||
POSIX:shlex.quote(sh 单引号语义)。
|
||||
Windows(run_command 经 cmd.exe / shell=True 执行):
|
||||
- 双引号包裹(& | < > ( ) ^ 在双引号内为字面量)
|
||||
- 含 " → 拒绝(闭合外层双引号,其后内容暴露给 cmd 解析)
|
||||
- 匹配 %VAR% 形态 → 拒绝(cmd 变量展开注入面)
|
||||
- 含换行/控制字符 → 拒绝(cmd 命令行不可含换行)
|
||||
"""
|
||||
value = str(value)
|
||||
if os.name != "nt":
|
||||
return shlex.quote(value), None
|
||||
if '"' in value:
|
||||
return "", f"❌ 参数含双引号,无法安全生成 Windows 命令: {value[:100]}"
|
||||
if "\n" in value or "\r" in value:
|
||||
return "", f"❌ 参数含换行符,无法生成命令: {value[:100]}"
|
||||
if _WIN_VAR_PATTERN.search(value):
|
||||
return "", (f"❌ 参数疑似含 cmd 变量引用(%变量名%),已拒绝以防命令注入: "
|
||||
f"{value[:100]}")
|
||||
return f'"{value}"', None
|
||||
|
||||
|
||||
def _ps_quote(value: str) -> str:
|
||||
"""PowerShell 单引号字面量(内部单引号翻倍),用于 -Command 内层参数。"""
|
||||
return "'" + str(value).replace("'", "''") + "'"
|
||||
|
||||
|
||||
# ==================== 文件操作工具基类 ====================
|
||||
|
||||
class FileToolBase(Tool):
|
||||
@@ -96,9 +132,64 @@ class ReadFileTool(FileToolBase):
|
||||
|
||||
# 限制max_lines上限,防止请求超大范围
|
||||
max_lines = min(max_lines, 500)
|
||||
quoted = shlex.quote(abs_path)
|
||||
# L1(审计 2026-08-18):start/end_line 为 LLM 可控参数,此前未强转
|
||||
# 直拼 sed 单引号模板——LLM 传含引号的字符串可闭合引号注入命令。
|
||||
# 强转 int(非法输入在此拦截)并夹取合理范围。
|
||||
try:
|
||||
start_line = int(start_line)
|
||||
end_line = int(end_line)
|
||||
except (TypeError, ValueError):
|
||||
return "❌ 参数错误: start_line/end_line 必须为整数"
|
||||
start_line = max(1, start_line)
|
||||
end_line = max(0, min(end_line, 10_000_000))
|
||||
|
||||
# 构造等价的命令行读取命令
|
||||
# T3(审计 2026-08-21):shlex.quote 是 POSIX 单引号语义,Windows 上
|
||||
# run_command 经 cmd.exe 执行——单引号不被识别,& | % < > 仍是元字符
|
||||
# (URL 含 & 即命令被截断;%VAR% 变量展开可注入)。按平台生成指引:
|
||||
# POSIX 保持 head/tail/sed,Windows 走 PowerShell 等价命令。
|
||||
quoted, qerr = _quote_arg(abs_path)
|
||||
if qerr:
|
||||
return qerr
|
||||
|
||||
if os.name == "nt":
|
||||
# PowerShell 等价命令(经 cmd 调 powershell,-Command 整体双引号
|
||||
# 包裹,内层路径用 _ps_quote 单引号字面量)
|
||||
if head:
|
||||
cmd = (f'powershell -NoProfile -Command "Get-Content -Path '
|
||||
f'{_ps_quote(abs_path)} -TotalCount {max_lines}"')
|
||||
wc_cmd = (f'powershell -NoProfile -Command "(Get-Content '
|
||||
f'{_ps_quote(abs_path)} | Measure-Object -Line).Lines"')
|
||||
elif tail:
|
||||
cmd = (f'powershell -NoProfile -Command "Get-Content -Path '
|
||||
f'{_ps_quote(abs_path)} -Tail {max_lines}"')
|
||||
wc_cmd = (f'powershell -NoProfile -Command "(Get-Content '
|
||||
f'{_ps_quote(abs_path)} | Measure-Object -Line).Lines"')
|
||||
else:
|
||||
if end_line > 0:
|
||||
first = start_line - 1
|
||||
take = end_line - start_line + 1
|
||||
elif max_lines > 0:
|
||||
first = start_line - 1
|
||||
take = max_lines
|
||||
else:
|
||||
first = start_line - 1
|
||||
take = 10_000_000
|
||||
cmd = (f'powershell -NoProfile -Command "Get-Content -Path '
|
||||
f'{_ps_quote(abs_path)} | Select-Object -Skip {first} '
|
||||
f'-First {take}"')
|
||||
wc_cmd = (f'powershell -NoProfile -Command "(Get-Content '
|
||||
f'{_ps_quote(abs_path)} | Measure-Object -Line).Lines"')
|
||||
prompt = f"""请使用 `run_command` 工具执行以下命令来读取文件内容,不要直接返回文件内容:
|
||||
|
||||
```
|
||||
{cmd}
|
||||
```
|
||||
|
||||
- 若需统计总行数:`{wc_cmd}`
|
||||
- 文件大小上限 {constants.MAX_FILE_SIZE} 字节,单次最多读取 500 行"""
|
||||
return prompt
|
||||
|
||||
# 构造等价的命令行读取命令(POSIX)
|
||||
if head:
|
||||
cmd = f"head -n {max_lines} {quoted}"
|
||||
elif tail:
|
||||
@@ -170,7 +261,40 @@ class WriteFileTool(FileToolBase):
|
||||
if not is_safe:
|
||||
return error
|
||||
|
||||
quoted = shlex.quote(abs_path)
|
||||
quoted, qerr = _quote_arg(abs_path)
|
||||
if qerr:
|
||||
return qerr
|
||||
|
||||
if os.name == "nt":
|
||||
# T3:Windows 下 heredoc 不可用(cmd 命令行不能含换行,且引号
|
||||
# 转义跨 cmd/PowerShell 两层极易出错)。content 整体 Base64 编码
|
||||
# 后经 PowerShell 解码写入——Base64 字符集 [A-Za-z0-9+/=] 不含
|
||||
# 任何 cmd/PowerShell 元字符,多行/引号/特殊字符内容全覆盖。
|
||||
b64 = base64.b64encode(content.encode("utf-8")).decode("ascii")
|
||||
if len(b64) > 6000:
|
||||
return ("❌ 写入内容过大:Windows 命令行长度上限约 8191 字符,"
|
||||
"Base64 编码后超出。请将内容分成多次调用(每次 ≤4KB),"
|
||||
"首次 mode=w、后续 mode=a 追加写入")
|
||||
writer = "Add-Content" if mode == "a" else "Set-Content"
|
||||
write_cmd = (f'powershell -NoProfile -Command "{writer} -Path '
|
||||
f'{_ps_quote(abs_path)} -Value ([Text.Encoding]::UTF8.'
|
||||
f'GetString([Convert]::FromBase64String(\'{b64}\')))"')
|
||||
directory = os.path.dirname(abs_path)
|
||||
if directory:
|
||||
write_cmd = (f'powershell -NoProfile -Command "New-Item '
|
||||
f'-ItemType Directory -Force -Path '
|
||||
f'{_ps_quote(directory)}" & {write_cmd}')
|
||||
prompt = f"""请使用 `run_command` 工具执行以下命令来写入文件,不要在此工具中直接写入:
|
||||
|
||||
```
|
||||
{write_cmd}
|
||||
```
|
||||
|
||||
- 写入内容已编码为 Base64(避免换行/引号/特殊字符的跨 cmd+PowerShell 转义问题),命令会自动解码并以 UTF-8 写入
|
||||
- 若需检查文件大小:`powershell -NoProfile -Command "(Get-Item {_ps_quote(abs_path)}).Length"`
|
||||
- 文件大小上限 {constants.MAX_FILE_SIZE} 字节"""
|
||||
return prompt
|
||||
|
||||
directory = os.path.dirname(abs_path)
|
||||
dir_cmd = f"mkdir -p {shlex.quote(directory)}" if directory else ""
|
||||
|
||||
@@ -244,7 +368,26 @@ class ListFilesTool(FileToolBase):
|
||||
if not os.path.isdir(abs_path):
|
||||
return f"❌ 路径不是目录: {abs_path}"
|
||||
|
||||
quoted = shlex.quote(abs_path)
|
||||
quoted, qerr = _quote_arg(abs_path)
|
||||
if qerr:
|
||||
return qerr
|
||||
|
||||
if os.name == "nt":
|
||||
# T3:Windows 走 PowerShell 等价命令(ls/du/find/file 无原生等价物)
|
||||
force_flag = " -Force" if show_hidden else ""
|
||||
cmd = (f'powershell -NoProfile -Command "Get-ChildItem{force_flag} '
|
||||
f'{_ps_quote(abs_path)} | Select-Object '
|
||||
f'Mode,LastWriteTime,Length,Name"')
|
||||
prompt = f"""请使用 `run_command` 工具执行以下命令来列出目录内容,不要直接在此工具中列出:
|
||||
|
||||
```
|
||||
{cmd}
|
||||
```
|
||||
|
||||
- 若需统计目录总大小:`powershell -NoProfile -Command "(Get-ChildItem -Recurse {_ps_quote(abs_path)} | Measure-Object Length -Sum).Sum"`
|
||||
- 若需递归列出:`powershell -NoProfile -Command "Get-ChildItem -Recurse -Depth 2 {_ps_quote(abs_path)} | Select-Object FullName"`"""
|
||||
return prompt
|
||||
|
||||
flag = "-la" if show_hidden else "-l"
|
||||
|
||||
# 返回 markdown 提示词,指导智能体用命令行工具列目录
|
||||
@@ -314,16 +457,58 @@ class SimpleWebCrawlerTool(Tool):
|
||||
host_port = proxy_url.split("://", 1)[-1]
|
||||
if "@" in host_port:
|
||||
host_port = host_port.split("@", 1)[-1]
|
||||
return f"--socks5 {shlex.quote(host_port)}"
|
||||
# T3:平台感知引用(原 shlex.quote 在 Windows cmd.exe 下无效)
|
||||
quoted, _err = _quote_arg(host_port)
|
||||
return f"--socks5 {quoted}" if quoted else ""
|
||||
|
||||
def execute(self, action: str, url: str, **kwargs) -> str:
|
||||
if action not in ("get", "text", "title", "links", "search"):
|
||||
return f"❌ 未知操作: {action}"
|
||||
|
||||
quoted_url = shlex.quote(url)
|
||||
# L10(审计 2026-08-18):url 为 LLM 可控输入(可被网页内容诱导),
|
||||
# 此前不限 scheme——file:// 可经生成的 curl 指引读取本地文件(SSRF)。
|
||||
# 白名单 http/https 后其余协议直接拒绝。
|
||||
if not url or not str(url).strip().lower().startswith(("http://", "https://")):
|
||||
return "❌ 仅支持 http/https 协议地址(已拦截 file:// 等非网络协议)"
|
||||
|
||||
quoted_url, qerr = _quote_arg(url)
|
||||
if qerr:
|
||||
return qerr
|
||||
timeout = constants.WEB_TIMEOUT
|
||||
user_agent = shlex.quote(constants.WEB_USER_AGENT)
|
||||
user_agent, ua_err = _quote_arg(constants.WEB_USER_AGENT)
|
||||
if ua_err:
|
||||
return ua_err
|
||||
proxy = self._proxy_curl_args()
|
||||
|
||||
if os.name == "nt":
|
||||
# T3:Windows 用系统内置 curl.exe(Win10 1803+);sed/grep 等
|
||||
# POSIX 管道无原生等价,分析类动作(text/title/links/search)
|
||||
# 指引先取原始 HTML 再由 AI 自行解析。URL 含 & 时靠双引号
|
||||
# 保护(cmd 双引号内为字面量),%VAR% 形态已在 _quote_arg 拒绝。
|
||||
base = f"curl.exe -sL --max-time {timeout} -A {user_agent} {proxy}".strip()
|
||||
if action == "get":
|
||||
cmd = f'{base} -w "\\nHTTP_CODE:%{{http_code}}" {quoted_url}'
|
||||
extra = "如需查看响应头:`curl.exe -sIL ...`;如需保存到文件:`curl.exe -sL ... -o 文件名`"
|
||||
elif action == "text":
|
||||
cmd = (f'{base} {quoted_url} | powershell -NoProfile -Command '
|
||||
f'"$input -replace \'<[^>]+>\',\'\'"')
|
||||
extra = "去除 HTML 标签提取纯文本"
|
||||
else:
|
||||
cmd = f"{base} {quoted_url}"
|
||||
extra = ("Windows 下无 grep/sed 管道,请取回原始 HTML 后自行分析"
|
||||
"(提取标题/链接/检索关键词)")
|
||||
prompt = f"""请使用 `run_command` 工具执行以下命令来完成网页操作,不要直接返回网页内容:
|
||||
|
||||
```
|
||||
{cmd}
|
||||
```
|
||||
|
||||
{extra}
|
||||
- 下载大小限制 {constants.WEB_MAX_DOWNLOAD_SIZE} 字节,可在命令后追加 `| powershell -NoProfile -Command "$input | Select-Object -First 100"` 限制
|
||||
- 当前已按配置自动附带代理:{proxy or '无(如需代理请先调用 set_proxy 工具)'}"""
|
||||
return prompt
|
||||
|
||||
# POSIX 分支:_quote_arg 在 POSIX 下即 shlex.quote,直接沿用
|
||||
base = f"curl -sL --max-time {timeout} -A {user_agent} {proxy}".strip()
|
||||
|
||||
if action == "get":
|
||||
@@ -342,7 +527,10 @@ class SimpleWebCrawlerTool(Tool):
|
||||
keyword = kwargs.get("keyword")
|
||||
if not keyword:
|
||||
return "❌ 需要提供 keyword 参数"
|
||||
cmd = f"{base} {quoted_url} | grep -c {shlex.quote(keyword)}"
|
||||
quoted_kw, kw_err = _quote_arg(keyword)
|
||||
if kw_err:
|
||||
return kw_err
|
||||
cmd = f"{base} {quoted_url} | grep -c {quoted_kw}"
|
||||
extra = f"返回匹配行数,大于 0 即表示找到关键词 '{keyword}'"
|
||||
|
||||
# 返回 markdown 提示词,指导智能体用 curl 完成网页访问
|
||||
|
||||
+14
-3
@@ -745,19 +745,30 @@ class MemoTool(Tool):
|
||||
return f"❌ 无法完成任务 #{memo_id}"
|
||||
|
||||
def _delete_memo(self, kwargs: Dict) -> str:
|
||||
"""删除任务"""
|
||||
"""删除任务(单条 memo_id 或批量 memo_ids)"""
|
||||
confirm = kwargs.get("confirm", False)
|
||||
if not confirm:
|
||||
return "❌ 需要设置 confirm=True 来确认删除操作"
|
||||
|
||||
memo_ids = kwargs.get("memo_ids")
|
||||
if memo_ids:
|
||||
# J10(审计 2026-08-19):schema 声明了 memo_ids 参数却从未实现,
|
||||
# 批量请求此前被静默忽略
|
||||
if not isinstance(memo_ids, (list, tuple)):
|
||||
return "❌ memo_ids 必须是列表"
|
||||
deleted = sum(1 for mid in memo_ids if self.memo_db.delete_memo(mid))
|
||||
return self.format_result(True, f"已删除 {deleted}/{len(memo_ids)} 条任务")
|
||||
|
||||
memo_id = kwargs.get("memo_id")
|
||||
if memo_id:
|
||||
if self.memo_db.complete_memo(memo_id, force_complete=True):
|
||||
# M5:真实物理删除(此前调 complete_memo 只标记完成,行不删除,
|
||||
# 重复任务被推进而非移除,stats 计数失真)
|
||||
if self.memo_db.delete_memo(memo_id):
|
||||
return self.format_result(True, f"任务 #{memo_id} 已删除")
|
||||
else:
|
||||
return f"❌ 无法删除任务 #{memo_id}"
|
||||
else:
|
||||
return "❌ 需要提供 memo_id"
|
||||
return "❌ 需要提供 memo_id 或 memo_ids"
|
||||
|
||||
def _search_memos(self, kwargs: Dict) -> str:
|
||||
"""搜索任务"""
|
||||
|
||||
@@ -39,10 +39,13 @@ class AskUserTool(Tool):
|
||||
)
|
||||
}
|
||||
},
|
||||
"required": []
|
||||
# L8:此前 required 为空,LLM 不传 question 时 str + None 抛 TypeError
|
||||
"required": ["question"]
|
||||
}
|
||||
|
||||
def execute(self, question: str, **kwargs) -> str:
|
||||
def execute(self, question: str = None, **kwargs) -> str:
|
||||
if not question:
|
||||
return "❌ 需要提供问题内容(question 参数)"
|
||||
return (
|
||||
"你遇到了想搞清楚的事情: " + question + "\n你现在可以询问用户。"
|
||||
)
|
||||
@@ -0,0 +1,424 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
tools.process_tools - 后台进程管理工具
|
||||
|
||||
五个 worker 执行面工具,配合 subsystems.process_manager.ProcessManager:
|
||||
- BgRunTool : bg_run 后台启动长驻进程,立即返回句柄
|
||||
- BgWriteTool : bg_write 向后台进程 stdin 写入(watermark 写读对齐)
|
||||
- BgListTool : bg_list 列出全部后台进程(运行中/已退出)
|
||||
- BgOutputTool : bg_output 拉取输出(tail 尾部 / since 游标增量)
|
||||
- BgKillTool : bg_kill 终止指定后台进程(关stdin→SIGTERM→SIGKILL)
|
||||
|
||||
设计语义(与用户约定的三决策):
|
||||
- 用户中断(Ctrl+C/取消事件)不影响后台进程
|
||||
- aibox 进程退出时默认终止全部,bg_run detach=true 豁免
|
||||
- 注册表全局共享:任何子代理可管理其他子代理启动的进程
|
||||
|
||||
交互语义:
|
||||
- stdin 为管道:行协议程序(REPL/问答 CLI)可 bg_write 写入;
|
||||
watermark = [stdin] 回显行序号,bg_output(since=watermark) 即本条
|
||||
输入的响应(写读对齐)
|
||||
- bg_run log_file=true 时输出 tee 到 workspace/bg_<句柄>.log(无界
|
||||
全档),可被原生工具(run_command tail/cat)读取
|
||||
|
||||
依赖:common (ConfigManager, constants) + tools.base.Tool +
|
||||
subsystems.process_manager(get_process_manager 单例)
|
||||
不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
from typing import Dict
|
||||
|
||||
from common import ConfigManager, constants
|
||||
from subsystems.process_manager import get_process_manager, ProcessManager
|
||||
from .base import Tool
|
||||
|
||||
|
||||
def _enabled() -> bool:
|
||||
"""后台进程功能开关:commands.enabled 与 background_processes.enabled
|
||||
双闸(后台执行本质仍是命令执行,前者关闭时一并拒绝)"""
|
||||
return constants.COMMANDS_ENABLED and constants.BG_PROCESSES_ENABLED
|
||||
|
||||
|
||||
def _fmt_duration(seconds: float) -> str:
|
||||
"""秒数 → 人读时长(如 3h05m / 12.5s)"""
|
||||
try:
|
||||
seconds = float(seconds)
|
||||
except (TypeError, ValueError):
|
||||
return "?"
|
||||
if seconds < 60:
|
||||
return f"{seconds:.1f}s"
|
||||
m, s = divmod(int(seconds), 60)
|
||||
if m < 60:
|
||||
return f"{m}m{s:02d}s"
|
||||
h, m = divmod(m, 60)
|
||||
return f"{h}h{m:02d}m{s:02d}s"
|
||||
|
||||
|
||||
def _coerce_int(value, default=None):
|
||||
"""命令面输入规整(W 批惯例):LLM 可能传字符串数字,统一安全转 int"""
|
||||
if value is None:
|
||||
return default
|
||||
try:
|
||||
return int(value)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
|
||||
def _coerce_bool(value, default: bool) -> bool:
|
||||
"""命令面输入规整:LLM 可能传字符串 "false"/"0",统一安全转 bool"""
|
||||
if value is None:
|
||||
return default
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
if isinstance(value, str):
|
||||
return value.strip().lower() not in ("false", "0", "no", "off", "")
|
||||
return bool(value)
|
||||
|
||||
|
||||
class BgRunTool(Tool):
|
||||
"""后台启动长驻进程"""
|
||||
|
||||
def __init__(self, manager: ProcessManager = None, logger=None, owner: str = None):
|
||||
super().__init__(logger)
|
||||
self.manager = manager or get_process_manager(logger)
|
||||
# T22:属主标识——写入/终止权限隔离(观察面 list/output 不限)
|
||||
self.owner = owner
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "bg_run"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return ("后台启动长驻进程(如 dev server、训练任务、watch 监控),立即返回句柄不阻塞。"
|
||||
"与 run_command 的区别:run_command 同步等待命令结束(超时即杀),适合短命令;"
|
||||
"bg_run 适合需要持续运行的进程,启动后用 bg_list 查看状态、bg_output 拉取输出、"
|
||||
"bg_kill 终止。stdin 为管道:行协议交互程序(REPL/问答式 CLI)可用 bg_write 写入。"
|
||||
"注意:用户中断不影响后台进程;aibox 退出时后台进程默认被终止,"
|
||||
"需要跨重启存活的进程用 detach=true(随后无法再经句柄管理)。")
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"command": {
|
||||
"type": "string",
|
||||
"description": "要执行的命令(shell 命令行)",
|
||||
},
|
||||
"proc_name": {
|
||||
"type": "string",
|
||||
"description": "进程标签(可选,便于识别,≤64 字符,默认取命令前 40 字符)",
|
||||
},
|
||||
"work_dir": {
|
||||
"type": "string",
|
||||
"description": "工作目录(默认 ~/.aibox/workspace)",
|
||||
},
|
||||
"env": {
|
||||
"type": "object",
|
||||
"description": "额外的环境变量",
|
||||
},
|
||||
"detach": {
|
||||
"type": "boolean",
|
||||
"description": "aibox 退出时不终止该进程(默认 false;detach 后进程脱离句柄管理)",
|
||||
},
|
||||
"log_file": {
|
||||
"type": "boolean",
|
||||
"description": "同时镜像全程输出到 workspace/bg_<句柄>.log(无界全档,可被 run_command 的 tail/cat 读取;默认 false 仅环形缓冲)",
|
||||
},
|
||||
},
|
||||
"required": ["command"],
|
||||
}
|
||||
|
||||
def execute(self, command: str = None, proc_name: str = None,
|
||||
work_dir: str = None, env: Dict = None,
|
||||
detach: bool = False, log_file: bool = False, **kwargs) -> str:
|
||||
if error := self.validate_params(["command"], command=command):
|
||||
return error
|
||||
if not _enabled():
|
||||
return "❌ 后台进程功能已禁用(commands.enabled 或 background_processes.enabled 为 false)"
|
||||
|
||||
# T18:LLM 可能传字符串 "false",裸 bool("false") is True——
|
||||
# detach 误开致进程脱离管理、log_file 误开,统一规整
|
||||
detach = _coerce_bool(detach, False)
|
||||
log_file = _coerce_bool(log_file, False)
|
||||
|
||||
result = self.manager.start(command, name=proc_name, work_dir=work_dir,
|
||||
env=env, detach=detach, log_file=log_file,
|
||||
owner=self.owner)
|
||||
if not result.get("ok"):
|
||||
return self.format_result(False, "后台进程启动失败",
|
||||
{"错误": result.get("error", "未知错误")})
|
||||
data = {
|
||||
"句柄": result["handle"],
|
||||
"PID": result["pid"],
|
||||
"名称": result["name"],
|
||||
"detach": "是(aibox 退出不终止)" if result.get("detach") else "否(aibox 退出时终止)",
|
||||
"后续": f"bg_output 拉取输出 / bg_list 查看状态 / bg_kill 终止(句柄 #{result['handle']})",
|
||||
}
|
||||
if result.get("log_file"):
|
||||
data["日志镜像"] = result["log_file"]
|
||||
return self.format_result(True, f"后台进程已启动(#{result['handle']})", data)
|
||||
|
||||
|
||||
class BgWriteTool(Tool):
|
||||
"""向后台进程 stdin 写入(交互式问答)"""
|
||||
|
||||
def __init__(self, manager: ProcessManager = None, logger=None, owner: str = None):
|
||||
super().__init__(logger)
|
||||
self.manager = manager or get_process_manager(logger)
|
||||
# T22:属主标识——仅属主可写(他人 stdin 不容注入)
|
||||
self.owner = owner
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "bg_write"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return ("向后台进程的标准输入写入一行(行协议交互:REPL / 问答式 CLI)。"
|
||||
"返回 watermark 水位:bg_output(handle=..., since=watermark) 得到的必然且仅是"
|
||||
"本条输入引发的进程响应(写读对齐,不含 [stdin] 回显行)。"
|
||||
"适用前提:进程按行读 stdin 并及时输出——建议启动时强制行缓冲"
|
||||
"(python -i -u、PYTHONUNBUFFERED=1、stdbuf -oL),否则响应可能滞留子进程缓冲。"
|
||||
"限制:无 TTY,全屏程序(vim/top 等)不可用;喂密码类交互(如 sudo -S)会留在"
|
||||
"对话记录中,敏感凭据请改用环境变量注入。")
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"handle": {
|
||||
"type": "integer",
|
||||
"description": "bg_run 返回的句柄编号",
|
||||
},
|
||||
"text": {
|
||||
"type": "string",
|
||||
"description": "要写入的文本(默认追加换行;多行输入会折叠为单行回显)",
|
||||
},
|
||||
"append_newline": {
|
||||
"type": "boolean",
|
||||
"description": "写入后是否自动补换行(默认 true;false 发送原始字节)",
|
||||
},
|
||||
"wait_lines": {
|
||||
"type": "integer",
|
||||
"description": "写完同步等待 N 行响应再返回(默认 0 立即返回,用 bg_output 轮询)",
|
||||
},
|
||||
"wait_timeout": {
|
||||
"type": "number",
|
||||
"description": "同步等待响应的超时秒数(默认 5,上限 60)",
|
||||
},
|
||||
},
|
||||
"required": ["handle", "text"],
|
||||
}
|
||||
|
||||
def execute(self, handle=None, text: str = None, append_newline=True,
|
||||
wait_lines=None, wait_timeout=None, **kwargs) -> str:
|
||||
if error := self.validate_params(["handle", "text"], handle=handle, text=text):
|
||||
return error
|
||||
if not _enabled():
|
||||
return "❌ 后台进程功能已禁用"
|
||||
|
||||
handle = _coerce_int(handle)
|
||||
if handle is None:
|
||||
return "❌ 错误:handle 必须是整数句柄"
|
||||
wait_i = _coerce_int(wait_lines, 0)
|
||||
wait_t = _coerce_int(wait_timeout, 5)
|
||||
|
||||
result = self.manager.write(
|
||||
handle, text,
|
||||
append_newline=_coerce_bool(append_newline, True),
|
||||
wait_lines=max(0, wait_i),
|
||||
wait_timeout=wait_t,
|
||||
owner=self.owner,
|
||||
)
|
||||
if not result.get("ok"):
|
||||
return self.format_result(False, f"写入后台进程 #{handle} 失败",
|
||||
{"错误": result.get("error", "未知错误")})
|
||||
|
||||
data = {
|
||||
"句柄": handle,
|
||||
"已写入": f"{result['sent_bytes']} 字节",
|
||||
"watermark": result["watermark"],
|
||||
"对齐读取": f"bg_output(handle={handle}, since={result['watermark']}) → 本条输入的响应",
|
||||
}
|
||||
if "wait" in result:
|
||||
w = result["wait"]
|
||||
data["等待输出"] = w.get("output") or "(超时未收到响应)"
|
||||
data["next_cursor"] = w.get("next_cursor")
|
||||
if w.get("timed_out"):
|
||||
data["警告"] = "等待响应超时(进程可能无输出或缓冲滞留),可用 bg_output 轮询"
|
||||
return self.format_result(True, f"已写入后台进程 #{handle}", data)
|
||||
|
||||
|
||||
class BgListTool(Tool):
|
||||
"""列出后台进程"""
|
||||
|
||||
def __init__(self, manager: ProcessManager = None, logger=None):
|
||||
super().__init__(logger)
|
||||
self.manager = manager or get_process_manager(logger)
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "bg_list"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return ("列出全部后台进程(含已退出的,保留尾部输出供验尸)。运行中在前。"
|
||||
"输出包含句柄/名称/PID/状态/运行时长,用句柄配合 bg_output/bg_kill 管理。")
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {"type": "object", "properties": {}}
|
||||
|
||||
def execute(self, **kwargs) -> str:
|
||||
if not _enabled():
|
||||
return "❌ 后台进程功能已禁用"
|
||||
procs = self.manager.list_procs()
|
||||
if not procs:
|
||||
return "✅ 当前没有后台进程(bg_run 启动)"
|
||||
lines = []
|
||||
status_emoji = {"running": "🟢", "exited": "⚪", "killed": "🔴", "expired": "⏹️"}
|
||||
for p in procs:
|
||||
emoji = status_emoji.get(p["status"], "❔")
|
||||
extra = ""
|
||||
if p["status"] == "running":
|
||||
extra = f"运行 {_fmt_duration(p['runtime_sec'])}"
|
||||
else:
|
||||
extra = f"已结束({_fmt_duration(p['runtime_sec'])},返回码 {p['exit_code']})"
|
||||
detach_tag = " [detach]" if p["detach"] else ""
|
||||
lines.append(f"{emoji} #{p['handle']} {p['name']}{detach_tag} "
|
||||
f"(pid {p['pid']}, {extra})")
|
||||
lines.append(f" 命令: {p['command']}")
|
||||
running = sum(1 for p in procs if p["status"] == "running")
|
||||
header = f"共 {len(procs)} 个后台进程(运行中 {running})"
|
||||
return "✅ " + header + "\n" + "\n".join(lines)
|
||||
|
||||
|
||||
class BgOutputTool(Tool):
|
||||
"""拉取后台进程输出"""
|
||||
|
||||
def __init__(self, manager: ProcessManager = None, logger=None):
|
||||
super().__init__(logger)
|
||||
self.manager = manager or get_process_manager(logger)
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "bg_output"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return ("拉取后台进程输出。两种模式:tail=N 取最近 N 行(默认 50);"
|
||||
"since=游标 增量拉取上次读取之后的新输出(游标用上次返回的 next_cursor)。"
|
||||
"长驻进程轮询建议用 since 模式避免重复传输旧输出。")
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"handle": {
|
||||
"type": "integer",
|
||||
"description": "bg_run 返回的句柄编号",
|
||||
},
|
||||
"tail": {
|
||||
"type": "integer",
|
||||
"description": "取最近 N 行(默认 50;与 since 二选一)",
|
||||
},
|
||||
"since": {
|
||||
"type": "integer",
|
||||
"description": "增量游标:上次返回的 next_cursor,取此序号之后的新行(与 tail 二选一)",
|
||||
},
|
||||
"max_bytes": {
|
||||
"type": "integer",
|
||||
"description": "本次拉取最大字节数(默认 64KB,上限 1MB)",
|
||||
},
|
||||
},
|
||||
"required": ["handle"],
|
||||
}
|
||||
|
||||
def execute(self, handle=None, tail=None, since=None, max_bytes=None, **kwargs) -> str:
|
||||
if error := self.validate_params(["handle"], handle=handle):
|
||||
return error
|
||||
if not _enabled():
|
||||
return "❌ 后台进程功能已禁用"
|
||||
|
||||
handle = _coerce_int(handle)
|
||||
if handle is None:
|
||||
return "❌ 错误:handle 必须是整数句柄"
|
||||
since_i = _coerce_int(since)
|
||||
tail_i = _coerce_int(tail)
|
||||
|
||||
result = self.manager.get_output(handle, tail=tail_i, since=since_i,
|
||||
max_bytes=_coerce_int(max_bytes, 65536))
|
||||
if not result.get("ok"):
|
||||
return self.format_result(False, "拉取输出失败",
|
||||
{"错误": result.get("error", "未知错误")})
|
||||
|
||||
data = {
|
||||
"状态": result["status"] + (f"(返回码 {result['exit_code']})"
|
||||
if result["exit_code"] is not None else ""),
|
||||
"本次行数": result["lines"],
|
||||
"next_cursor": result["next_cursor"],
|
||||
}
|
||||
if result.get("truncated"):
|
||||
data["警告"] = "输出超过本次字节上限,已截断(提高 max_bytes 或用 since 分段拉取)"
|
||||
if result.get("evicted"):
|
||||
data["提示"] = (f"有 {result['evicted']} 行早期输出因环形缓冲淘汰不可读取"
|
||||
"(需要完整日志请让进程重定向到文件)")
|
||||
output = result.get("output") or "(本次无新输出)"
|
||||
data["输出"] = output
|
||||
return self.format_result(True, f"后台进程 #{handle} 输出", data)
|
||||
|
||||
|
||||
class BgKillTool(Tool):
|
||||
"""终止后台进程"""
|
||||
|
||||
def __init__(self, manager: ProcessManager = None, logger=None, owner: str = None):
|
||||
super().__init__(logger)
|
||||
self.manager = manager or get_process_manager(logger)
|
||||
# T22:属主标识——仅属主可终止(防提示注入跨会话杀 dev server)
|
||||
self.owner = owner
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "bg_kill"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return ("终止后台进程(整个进程组,含子进程)。默认先 SIGTERM 优雅终止,"
|
||||
"宽限期内未退出再 SIGKILL;force=true 直接 SIGKILL。对已退出进程幂等。")
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"handle": {
|
||||
"type": "integer",
|
||||
"description": "bg_run 返回的句柄编号",
|
||||
},
|
||||
"force": {
|
||||
"type": "boolean",
|
||||
"description": "跳过优雅终止直接 SIGKILL(默认 false)",
|
||||
},
|
||||
"grace": {
|
||||
"type": "number",
|
||||
"description": "SIGTERM 后等待进程退出的宽限秒数(默认 3,上限 30)",
|
||||
},
|
||||
},
|
||||
"required": ["handle"],
|
||||
}
|
||||
|
||||
def execute(self, handle=None, force: bool = False, grace: float = 3.0, **kwargs) -> str:
|
||||
if error := self.validate_params(["handle"], handle=handle):
|
||||
return error
|
||||
if not _enabled():
|
||||
return "❌ 后台进程功能已禁用"
|
||||
|
||||
handle = _coerce_int(handle)
|
||||
if handle is None:
|
||||
return "❌ 错误:handle 必须是整数句柄"
|
||||
|
||||
# T18:force 字符串 "false" 会被裸 bool 判为 True,跳过宽限直接 SIGKILL
|
||||
force = _coerce_bool(force, False)
|
||||
|
||||
result = self.manager.kill(handle, force=force, grace=grace, owner=self.owner)
|
||||
if not result.get("ok"):
|
||||
return self.format_result(False, f"终止后台进程 #{handle} 失败",
|
||||
{"错误": result.get("error", "未知错误")})
|
||||
data = {
|
||||
"句柄": handle,
|
||||
"状态": result["status"],
|
||||
}
|
||||
if result.get("exit_code") is not None:
|
||||
data["返回码"] = result["exit_code"]
|
||||
if result.get("already_finished"):
|
||||
return self.format_result(True, f"后台进程 #{handle} 此前已结束", data)
|
||||
return self.format_result(True, f"后台进程 #{handle} 已终止", data)
|
||||
@@ -56,6 +56,11 @@ class UpdateIdentityTool(Tool):
|
||||
if not content:
|
||||
return "❌ 需要提供身份信息内容"
|
||||
|
||||
# N9:持久化注入面约束——内容超限拒绝写入(管理端另有兜底)
|
||||
if len(content) > PromptManager.MAX_CONTENT_LENGTH:
|
||||
return (f"❌ 内容过长({len(content)} > {PromptManager.MAX_CONTENT_LENGTH} 字符),"
|
||||
f"身份文件作为系统提示词加载,请精简后重试")
|
||||
|
||||
if not confirm:
|
||||
return "❌ 身份更新已取消,如需更新请设置 confirm=true"
|
||||
|
||||
@@ -119,6 +124,11 @@ class UpdateSoulTool(Tool):
|
||||
if not content:
|
||||
return "❌ 需要提供智能体自我认识定义内容"
|
||||
|
||||
# N9:持久化注入面约束——内容超限拒绝写入(管理端另有兜底)
|
||||
if len(content) > PromptManager.MAX_CONTENT_LENGTH:
|
||||
return (f"❌ 内容过长({len(content)} > {PromptManager.MAX_CONTENT_LENGTH} 字符),"
|
||||
f"灵魂文件作为系统提示词加载,请精简后重试")
|
||||
|
||||
if not confirm:
|
||||
return "❌ 自我认识更新已取消,如需更新请设置 confirm=true"
|
||||
|
||||
@@ -223,6 +233,11 @@ class SaveMemoryAndEndConversationTool(Tool):
|
||||
status_content += f"【下一步计划】\n{next_goals}\n\n"
|
||||
status_content += f"【对话总结】\n{summary}"
|
||||
|
||||
# N9:状态文件作为系统提示词加载,超长拒绝写入(约束注入规模)
|
||||
if len(status_content) > PromptManager.MAX_CONTENT_LENGTH:
|
||||
return (f"❌ 状态内容过长({len(status_content)} > "
|
||||
f"{PromptManager.MAX_CONTENT_LENGTH} 字符),请精简 current_status/summary 后重试")
|
||||
|
||||
# 更新状态文件
|
||||
status_updated = self.prompt_manager.update_status(status_content)
|
||||
|
||||
|
||||
@@ -146,7 +146,7 @@ class SetProxyTool(Tool):
|
||||
check_detail = detail
|
||||
|
||||
config = ConfigManager()
|
||||
config.set_proxy(host, port, username, password, password_env, enabled, save)
|
||||
saved = config.set_proxy(host, port, username, password, password_env, enabled, save)
|
||||
proxy_info = config.get_proxy_info()
|
||||
if proxy_info.get("password_env"):
|
||||
pw_source = "环境变量 " + proxy_info["password_env"]
|
||||
@@ -159,10 +159,13 @@ class SetProxyTool(Tool):
|
||||
"服务器": f"{proxy_info['host']}:{proxy_info['port']}",
|
||||
"认证": "有" if proxy_info.get("username") else "无",
|
||||
"密码来源": pw_source,
|
||||
"持久化": "已保存" if save else "仅内存"
|
||||
"持久化": "已保存" if saved else "保存失败(仅本次会话有效,重启后丢失)"
|
||||
}
|
||||
if check_detail:
|
||||
data["连通性"] = check_detail
|
||||
# L2:落盘失败(磁盘只读等)如实报告,不再静默报"已保存"
|
||||
if save and not saved:
|
||||
return self.format_result(False, "代理已配置但保存失败(仅本次会话有效,重启后丢失)", data)
|
||||
return self.format_result(True, "代理设置成功", data)
|
||||
|
||||
|
||||
|
||||
+37
-14
@@ -16,13 +16,18 @@ from common import Logger, constants, i18n
|
||||
from .base import Tool
|
||||
|
||||
# 工具参数名命中这些关键词时,debug 日志里打码(防密码/密钥进日志文件)
|
||||
_SENSITIVE_PARAM_KEYS = {"password", "api_key", "token", "secret", "authorization"}
|
||||
# L11(审计 2026-08-18):精确匹配漏掉 apiKey/access_token/passwd 等变体,
|
||||
# 改包含匹配——日志场景下过度打码可接受,漏打码才是问题。
|
||||
_SENSITIVE_PARAM_KEYS = (
|
||||
"password", "passwd", "api_key", "apikey", "token", "secret",
|
||||
"authorization", "credential",
|
||||
)
|
||||
|
||||
|
||||
def _sanitize_kwargs(kwargs: Dict) -> Dict:
|
||||
"""复制一份参数并把敏感键的值替换为 ***(仅用于日志输出)"""
|
||||
return {
|
||||
k: ("***" if str(k).lower() in _SENSITIVE_PARAM_KEYS else v)
|
||||
k: ("***" if any(s in str(k).lower() for s in _SENSITIVE_PARAM_KEYS) else v)
|
||||
for k, v in kwargs.items()
|
||||
}
|
||||
|
||||
@@ -63,12 +68,22 @@ class ToolRegistry:
|
||||
"""获取所有工具的OpenAI函数调用格式"""
|
||||
return [tool.to_function_call_format() for tool in self._tools.values()]
|
||||
|
||||
@staticmethod
|
||||
def _truncate_result(result) -> str:
|
||||
"""T15:工具返回值统一长度预算(非流式与流式最终结果同口径)"""
|
||||
if not isinstance(result, str):
|
||||
return result
|
||||
limit = constants.TOOL_MAX_RESULT_LENGTH
|
||||
if len(result) > limit:
|
||||
return (result[:limit] +
|
||||
f"\n⚠️ 结果超过 {limit} 字符上限,已截断")
|
||||
return result
|
||||
|
||||
def set_output_queue(self, queue: queue.Queue):
|
||||
"""设置输出队列"""
|
||||
"""设置输出队列(T4:仅记注册表级引用,不再传播到共享工具实例——
|
||||
实例属性由 execute_streaming 每次调用自行注入局部队列,全局传播
|
||||
会让并发会话互相覆盖对方的输出通路,造成跨会话输出串扰)"""
|
||||
self.current_output_queue = queue
|
||||
# 同时设置所有工具的队列
|
||||
for tool in self._tools.values():
|
||||
tool.output_queue = queue
|
||||
|
||||
def execute_tool(self, name: str, **kwargs) -> str:
|
||||
"""执行工具"""
|
||||
@@ -80,11 +95,14 @@ class ToolRegistry:
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# T4:入口快照会话队列。转发循环可能持续数秒~小时,期间并发
|
||||
# 会话的 set_output_queue 会覆盖单值引用——全程用快照,保证
|
||||
# 本会话输出只进本会话队列,结束清理也不影响在途转发
|
||||
session_queue = self.current_output_queue
|
||||
# 检查是否有队列可用,并且工具支持流式输出
|
||||
if self.current_output_queue and hasattr(tool, 'execute_streaming'):
|
||||
if session_queue and hasattr(tool, 'execute_streaming'):
|
||||
# 设置工具的流式输出属性
|
||||
tool.streaming_output = True
|
||||
tool.output_queue = self.current_output_queue
|
||||
|
||||
# 创建队列来收集流式输出
|
||||
output_queue = queue.Queue()
|
||||
@@ -129,8 +147,8 @@ class ToolRegistry:
|
||||
# 预算耗尽:停止转发,仅提示一次
|
||||
if not truncation_noted:
|
||||
truncation_noted = True
|
||||
if self.current_output_queue:
|
||||
self.current_output_queue.put(
|
||||
if session_queue:
|
||||
session_queue.put(
|
||||
("line", f"\n⚠️ 流式输出超过 {line_budget} 字符上限,已停止转发剩余内容\n"))
|
||||
continue
|
||||
remaining = line_budget - forwarded
|
||||
@@ -141,13 +159,13 @@ class ToolRegistry:
|
||||
else:
|
||||
forwarded += len(piece)
|
||||
# 发送到当前队列(会通过WebSocket发送)
|
||||
if self.current_output_queue:
|
||||
self.current_output_queue.put(("line", piece))
|
||||
if session_queue:
|
||||
session_queue.put(("line", piece))
|
||||
elif msg_type == "complete":
|
||||
final_result = content
|
||||
# 也发送完整结果
|
||||
if self.current_output_queue:
|
||||
self.current_output_queue.put(("complete", content))
|
||||
if session_queue:
|
||||
session_queue.put(("complete", content))
|
||||
except queue.Empty:
|
||||
continue
|
||||
|
||||
@@ -169,6 +187,11 @@ class ToolRegistry:
|
||||
# 普通执行
|
||||
result = tool.execute(**kwargs)
|
||||
|
||||
# T15:两条分支的最终返回值统一过 TOOL_MAX_RESULT_LENGTH——
|
||||
# 此前仅流式转发有 line_budget(N8),非流式结果无上限直入
|
||||
# 对话上下文(memory/memo get 超长 content,token/内存放大)
|
||||
result = self._truncate_result(result)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
self.logger.debug(f"工具结果 ({elapsed:.2f}s): {result[:100]}...")
|
||||
return result
|
||||
|
||||
@@ -16,6 +16,7 @@ tools.skill_tools - 技能相关工具
|
||||
import os
|
||||
import queue
|
||||
import re
|
||||
import shlex
|
||||
from typing import Dict, Optional
|
||||
|
||||
from common import constants, i18n, Logger
|
||||
@@ -253,10 +254,14 @@ class InstallSkillTool(Tool):
|
||||
skills_dir = constants.SKILLS_DIR
|
||||
|
||||
# 如果给了具体路径,在说明里带入,否则用占位符
|
||||
src = skill_path.strip() if skill_path else "/path/to/your_skill.md"
|
||||
raw_src = skill_path.strip() if skill_path else ""
|
||||
# 去掉扩展名作为技能名称,再统一补 .md 作为目标文件名
|
||||
base = os.path.splitext(os.path.basename(src))[0] if src else "your_skill"
|
||||
base = os.path.splitext(os.path.basename(raw_src))[0] if raw_src else "your_skill"
|
||||
dst = os.path.join(skills_dir, base + ".md")
|
||||
# L7(审计 2026-08-18):src 为 LLM 可控输入,含空格/元字符时
|
||||
# 未经引用的命令会被 shell 拆词/注入(指引由子代理经 run_command 执行)
|
||||
src = shlex.quote(raw_src) if raw_src else "/path/to/your_skill.md"
|
||||
dst_q = shlex.quote(dst)
|
||||
|
||||
guide = f"""📦 扩展技能安装指引
|
||||
{"=" * 55}
|
||||
@@ -282,7 +287,7 @@ class InstallSkillTool(Tool):
|
||||
|
||||
调用 create_subagent 派生子代理执行复制命令(目标文件名统一补 .md 扩展名):
|
||||
|
||||
cp {src} {dst}
|
||||
cp {src} {dst_q}
|
||||
|
||||
⚠️ 注意:目标文件名必须带 .md 扩展名,
|
||||
系统以 .md 文件作为技能加载,且技能名称取文件名(不含 .md)。
|
||||
|
||||
+61
-12
@@ -181,6 +181,9 @@ class WebSocketClient:
|
||||
self.connected = False
|
||||
print(f"\n{Colors.ERROR}❌ 连接已断开{Colors.RESET}")
|
||||
except Exception as e:
|
||||
# J14:异常退出同样标记断开(此前 connected 恒 True,用户继续
|
||||
# "成功"发送但永远收不到回复,成僵尸连接)
|
||||
self.connected = False
|
||||
print(f"\n{Colors.ERROR}❌ 接收消息错误: {e}{Colors.RESET}")
|
||||
|
||||
async def handle_message(self, data: dict, tts: TextToSpeech):
|
||||
@@ -266,12 +269,13 @@ class WebSocketClient:
|
||||
tts: TextToSpeech = None, task_info: dict = None):
|
||||
if msg_type == 'chat':
|
||||
stream = self.streaming_messages['chat']
|
||||
stream.add_chunk(content)
|
||||
|
||||
# J14(审计 2026-08-19):前缀判断必须在 add_chunk 之前——
|
||||
# add_chunk 会置 is_active=True,原顺序使 "AI助手:" 前缀永不打印
|
||||
if not stream.is_active:
|
||||
print(f"\n{Colors.AI}🤖 AI助手: ", end='', flush=True)
|
||||
stream.is_active = True
|
||||
|
||||
stream.add_chunk(content)
|
||||
print(content, end='', flush=True)
|
||||
|
||||
elif msg_type == 'task' and task_info:
|
||||
@@ -343,12 +347,17 @@ class VoiceInterface:
|
||||
with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as f:
|
||||
test_file = f.name
|
||||
|
||||
result = subprocess.run(
|
||||
['arecord', '-f', 'cd', '-d', '1', test_file],
|
||||
capture_output=True, text=True
|
||||
)
|
||||
|
||||
os.unlink(test_file)
|
||||
try:
|
||||
result = subprocess.run(
|
||||
['arecord', '-f', 'cd', '-d', '1', test_file],
|
||||
capture_output=True, text=True,
|
||||
timeout=5, # M7(审计 2026-08-20):设备挂死此前永久卡死
|
||||
)
|
||||
finally:
|
||||
# M7:超时/异常路径同样清理探测临时文件(此前残留 /tmp);
|
||||
# TimeoutExpired 传播至外层 except,按"麦克风不可用"降级
|
||||
if os.path.exists(test_file):
|
||||
os.unlink(test_file)
|
||||
|
||||
if result.returncode == 0:
|
||||
print(f"{Colors.INFO}🎤 arecord 检测到设备,将使用 arecord 录音{Colors.RESET}")
|
||||
@@ -428,6 +437,7 @@ class VoiceInterface:
|
||||
return None
|
||||
|
||||
print(f"{Colors.INFO}🎤 正在听... (按 Ctrl+C 取消){Colors.RESET}")
|
||||
temp_file = None # J14:录音临时文件,任何异常路径统一清理
|
||||
|
||||
try:
|
||||
if self.use_arecord:
|
||||
@@ -506,6 +516,13 @@ class VoiceInterface:
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f"{Colors.ERROR}❌ 语音输入错误: {e}{Colors.RESET}")
|
||||
# J14:录音临时文件(含用户语音)任何异常路径都清理——
|
||||
# 此前 wave.open/whisper 解析失败、arecord 超时等滞留 /tmp
|
||||
if temp_file and os.path.exists(temp_file):
|
||||
try:
|
||||
os.unlink(temp_file)
|
||||
except OSError:
|
||||
pass
|
||||
return None
|
||||
|
||||
async def handle_command(self, command: str):
|
||||
@@ -602,17 +619,49 @@ class VoiceInterface:
|
||||
print(f"\n{Colors.USER}👤 用户: {message}{Colors.RESET}")
|
||||
await self.client.send_message(message)
|
||||
|
||||
def _stdin_reader(self, loop: asyncio.AbstractEventLoop):
|
||||
"""M1(审计 2026-08-20):单守护线程读 stdin——替代 asyncio.to_thread
|
||||
每行一个 input 线程的方案:Ctrl+C 取消 await 后旧线程仍阻塞 readline
|
||||
永久泄漏,并"偷吃"后续输入整行(表现为偶发丢行)。
|
||||
方案照搬 client.py:735-755(守护线程 + Queue + call_soon_threadsafe)。"""
|
||||
def _put(item):
|
||||
try:
|
||||
loop.call_soon_threadsafe(self._input_queue.put_nowait, item)
|
||||
except RuntimeError:
|
||||
pass # 事件循环已关闭(程序退出中)
|
||||
|
||||
while self._running:
|
||||
try:
|
||||
line = sys.stdin.readline()
|
||||
except Exception:
|
||||
break
|
||||
if not line: # EOF (Ctrl+D)
|
||||
_put(None)
|
||||
break
|
||||
_put(line)
|
||||
|
||||
async def run(self):
|
||||
self.print_banner()
|
||||
|
||||
prompt = f"{Colors.USER}>>> {Colors.RESET}"
|
||||
|
||||
# M1:进程生命周期只启动一次 stdin 守护线程
|
||||
self._input_queue = asyncio.Queue()
|
||||
self._running = True
|
||||
loop = asyncio.get_running_loop()
|
||||
if getattr(self, '_stdin_thread', None) is None:
|
||||
self._stdin_thread = threading.Thread(
|
||||
target=self._stdin_reader, args=(loop,),
|
||||
daemon=True, name="stdin-reader")
|
||||
self._stdin_thread.start()
|
||||
|
||||
while True:
|
||||
try:
|
||||
user_input = await asyncio.to_thread(
|
||||
lambda: input(prompt)
|
||||
)
|
||||
user_input = user_input.strip()
|
||||
print(prompt, end='', flush=True)
|
||||
line = await self._input_queue.get()
|
||||
if line is None: # EOF (Ctrl+D)
|
||||
raise EOFError
|
||||
user_input = line.strip()
|
||||
|
||||
if not user_input:
|
||||
if self.has_microphone:
|
||||
|
||||
@@ -21,6 +21,7 @@
|
||||
spi4cj = "1.0.0"
|
||||
#lua4cj = { git = "https://gitcode.com/qingss0/lua4cj.git", branch="main" }
|
||||
lua4cj = "1.0.0"
|
||||
nn4cj = { git = "https://gitcode.com/qingss0/nn4cj.git", branch="main"}
|
||||
# 这种x86机器人开发板感觉不经常使用,我没用过,不知道里面的i2c调用是不是参数是一样的
|
||||
[target.x86_64-unknown-linux-gnu]
|
||||
[target.x86_64-unknown-linux-gnu.dependencies]
|
||||
|
||||
+7
-7
@@ -98,13 +98,13 @@ class MatrixObject {
|
||||
}
|
||||
|
||||
/* ---------------- 对Nn1进行扩展 ---------------- */
|
||||
extend Nn1 {
|
||||
/* -------- 转换为MatrixObject对象 -------- */
|
||||
public func toMto(): MatrixObject {
|
||||
let mto = MatrixObject()
|
||||
return mto
|
||||
}
|
||||
}
|
||||
//extend Nn1 {
|
||||
// /* -------- 转换为MatrixObject对象 -------- */
|
||||
// public func toMto(): MatrixObject {
|
||||
// let mto = MatrixObject()
|
||||
// return mto
|
||||
// }
|
||||
//}
|
||||
|
||||
/* -------- 矩阵式解析器 --------
|
||||
* TODO(20260514-qing): 这个类输入流数据,输出MatrixObject对象
|
||||
|
||||
@@ -1,152 +0,0 @@
|
||||
/*
|
||||
* i2c - 用户界面
|
||||
*
|
||||
* 这里已经提供了I2C通用接口,可以将上层代码移动到仓颉下面
|
||||
*/
|
||||
|
||||
#include <stdio.h>
|
||||
#include <stdint.h>
|
||||
#include <fcntl.h>
|
||||
#ifdef __linux__
|
||||
#include <unistd.h>
|
||||
#include <sys/ioctl.h>
|
||||
#include <linux/i2c-dev.h> /* 其中宏的取值可能具有平台特异性 */
|
||||
#include <i2c/smbus.h>
|
||||
#endif
|
||||
|
||||
#include "bot4cj.h" /* 有没有可能在仓颉代码中解析C头文件取得宏 */
|
||||
|
||||
|
||||
/*
|
||||
* 初始化i2c设备
|
||||
*
|
||||
* 动态指定设备文件
|
||||
*
|
||||
* 返回打开的文件句柄或者错误码
|
||||
*/
|
||||
#ifdef __linux__
|
||||
int16_t I2C_init(const char* filepath) {
|
||||
int16_t fd;
|
||||
if ((fd = open(filepath, O_RDWR)) < 0) {
|
||||
perror("Failed to open I2C bus");
|
||||
return BOT4CJ_I2C_OPENED_FAILED;
|
||||
}
|
||||
else {
|
||||
return fd;
|
||||
}
|
||||
}
|
||||
#endif
|
||||
#ifdef _WIN64
|
||||
int16_t I2C_init(const char* filepath) {
|
||||
return BOT4CJ_I2C_OPENED_FAILED;
|
||||
}
|
||||
#endif
|
||||
|
||||
/* 释放I2C设备文件 */
|
||||
#ifdef __linux__
|
||||
void I2C_free(int16_t fd) {
|
||||
close(fd);
|
||||
}
|
||||
#endif
|
||||
#ifdef _WIN64
|
||||
void I2C_free(int16_t fd) {
|
||||
/* Windows平台下无需操作 */
|
||||
}
|
||||
#endif
|
||||
|
||||
|
||||
|
||||
/* 设置I2C地址 */
|
||||
#ifdef __linux__
|
||||
int I2C_setAddr(int16_t fd, uint8_t addr) {
|
||||
if (ioctl(fd, I2C_SLAVE, addr) < 0) {
|
||||
perror("Failed to set I2C address");
|
||||
return BOT4CJ_I2C_ERR;
|
||||
}
|
||||
else {
|
||||
return BOT4CJ_I2C_OK;
|
||||
}
|
||||
}
|
||||
#endif
|
||||
#ifdef _WIN64
|
||||
int I2C_setAddr(int16_t fd, uint8_t addr) {
|
||||
return BOT4CJ_I2C_ERR;
|
||||
}
|
||||
#endif
|
||||
|
||||
/* 写寄存器 */
|
||||
#ifdef __linux__
|
||||
int I2C_writeReg(int16_t fd, uint8_t reg, uint8_t val) {
|
||||
/*uint8_t buf[] = {reg, val};
|
||||
if (write(fd, buf, 2) != 2) {
|
||||
perror("I2C write error");
|
||||
return BOT4CJ_I2C_ERR;
|
||||
}
|
||||
else {
|
||||
return BOT4CJ_I2C_OK;
|
||||
}*/
|
||||
|
||||
return i2c_smbus_write_byte_data(fd, reg, val);
|
||||
}
|
||||
#endif
|
||||
#ifdef _WIN64
|
||||
int I2C_writeReg(int16_t fd, uint8_t reg, uint8_t val) {
|
||||
return BOT4CJ_I2C_ERR;
|
||||
}
|
||||
#endif
|
||||
|
||||
/* 读寄存器 失败返回0 */
|
||||
#ifdef __linux__
|
||||
uint8_t I2C_readReg(int16_t fd, uint8_t reg) {
|
||||
|
||||
uint8_t value;
|
||||
|
||||
/*if (write(fd, ®, 1) != 1) {
|
||||
perror("I2C read error");
|
||||
return 0;
|
||||
}
|
||||
|
||||
if (read(fd, &value, 1)!= 1) {
|
||||
perror("I2C read error");
|
||||
return 0;
|
||||
}*/
|
||||
|
||||
//return value;
|
||||
|
||||
int8_t a = i2c_smbus_read_byte_data(fd, reg);
|
||||
return *(uint8_t *)& a;
|
||||
}
|
||||
#endif
|
||||
#ifdef _WIN64
|
||||
uint8_t I2C_readReg(int16_t fd, uint8_t reg) {
|
||||
return 0;
|
||||
}
|
||||
#endif
|
||||
|
||||
/* 读寄存器 - 有符号
|
||||
* 失败返回0 */
|
||||
#ifdef __linux__
|
||||
int8_t I2C_readRegFlag(int16_t fd, uint8_t reg) {
|
||||
|
||||
int8_t value;
|
||||
|
||||
//if (write(fd, ®, 1) != 1) {
|
||||
//perror("I2C read error");
|
||||
//return 0;
|
||||
//}
|
||||
|
||||
//if (read(fd, &value, 1)!= 1) {
|
||||
//perror("I2C read error");
|
||||
//return 0;
|
||||
//}
|
||||
|
||||
//return value;
|
||||
|
||||
return i2c_smbus_read_byte_data(fd, reg);
|
||||
}
|
||||
#endif
|
||||
#ifdef _WIN64
|
||||
int8_t I2C_readRegFlag(int16_t fd, uint8_t reg) {
|
||||
return 0;
|
||||
}
|
||||
#endif
|
||||
@@ -1,516 +0,0 @@
|
||||
#include <random>
|
||||
#include <algorithm>
|
||||
#include "Matrix.hh"
|
||||
|
||||
namespace qing {
|
||||
|
||||
template <typename T>
|
||||
Matrx<T>::Matrx(): row(0), col(0) {}
|
||||
|
||||
template <typename T>
|
||||
Matrx<T>::Matrx(const uint64_t row, const uint64_t col)
|
||||
: row(row), col(col) { /* 该构造函数用于从头开始构造 */
|
||||
arr = std::vector<T>(col * row);
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
Matrx<T>::Matrx(const uint64_t row, const uint64_t col, const std::vector<T>& src)
|
||||
: row(row), col(col) { /* 该构造函数用于通过已有的vector进行构造 */
|
||||
arr = src;
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
const T& Matrx<T>::get(uint64_t row, uint64_t col) const { /* 访问器 - 访问矩阵元素 */
|
||||
return arr[col * this->row + row];
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
T& Matrx<T>::witch(uint64_t row, uint64_t col) { /* 访问器 - 访问矩阵元素 */
|
||||
return arr[col * this->row + row];
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
void Matrx<T>::set(uint64_t row, uint64_t col, T v) { /* 设置矩阵元素 */
|
||||
arr[col * this->row + row] = v;
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
const std::vector<T>& Matrx<T>::get_data() const{ /* 访问器 - 用于获取原始数据 */
|
||||
return arr;
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
const uint64_t Matrx<T>::get_col() const { /* 获取列数 */
|
||||
return col;
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
const uint64_t Matrx<T>::get_row() const { /* 获取行数 */
|
||||
return row;
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
bool Matrx<T>::has_same_shape(const Matrx<T>& other) { /* 两个矩阵是否具有相同形状 */
|
||||
return other.get_col() == this->get_col() && other.get_row() == this->get_row();
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
bool Matrx<T>::has_dot_shape(const Matrx<T>& other) { /* 具有矩阵乘法的形状 */
|
||||
return this->get_col() == other.get_row();
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
void Matrx<T>::assert_same_shape(const Matrx<T>& other) { /* 断言两个矩阵具有相同形状 */
|
||||
if (!has_same_shape(other))
|
||||
throw std::invalid_argument("The shapes of matrix are not same.");
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
void Matrx<T>::assert_dot_shape(const Matrx<T>& other) { /* 断言两个矩阵具有点积前置形状 */
|
||||
if (!has_dot_shape(other)) {
|
||||
throw std::invalid_argument("The shapes of matrix can not be dot.");
|
||||
}
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
Matrx<T> Matrx<T>::operator+(const Matrx<T>& other) { /* 加法操作符重定义 - 矩阵加法 */
|
||||
//assert_same_shape(other); /* 检查两个矩阵形状是否一致 */
|
||||
std::vector<T> v0 = this->get_data();
|
||||
std::vector<T> vec = other.get_data();
|
||||
for (auto it = v0.begin(), it_or = vec.begin(); it != v0.end(); ++it, ++it_or )
|
||||
*it += *it_or;
|
||||
return Matrx<T>{this->get_row(), this->get_col(), v0};
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
Matrx<T> Matrx<T>::operator+(T& item) { /* 矩阵加法广播 */
|
||||
auto len = this->get_data().size();
|
||||
auto v0 = std::vector<T> (len, item);
|
||||
auto matrix = (*this) + Matrx(get_row(), get_col(), v0);
|
||||
|
||||
//auto matrix = *this;
|
||||
//for (auto it = arr.begin(); it != arr.end(); ++it)
|
||||
// *it += item;
|
||||
return matrix;
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
Matrx<T> Matrx<T>::operator-(const Matrx<T>& other) { /* 操作符重载 - 矩阵减法 */
|
||||
//assert_same_shape(other); /* 检查两个矩阵形状是否一致 */
|
||||
std::vector<T> v0 = this->get_data();
|
||||
std::vector <T> vec = other.get_data();
|
||||
for (auto it = v0.begin(), it_or = vec.begin(); it != v0.end(); ++it, ++it_or )
|
||||
*it -= *it_or;
|
||||
return Matrx{this->get_row(), this->get_col(), v0};
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
Matrx<T> Matrx<T>::operator-(T& item) { /* 操作符重载 - 矩阵减法广播 */
|
||||
auto len = this->get_data().size();
|
||||
auto v0 = std::vector<T> (len, item);
|
||||
auto matrix = (*this) - Matrx(get_row(), get_col(), v0);
|
||||
|
||||
//Matrx<T> matrix = *this;
|
||||
//for (auto it = arr.begin(); it != arr.end(); ++it)
|
||||
// *it -= item;
|
||||
return matrix;
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
Matrx<T> Matrx<T>::dot(const Matrx<T>& other) { /* 操作符重载 - 矩阵乘法 */
|
||||
//assert_dot_shape(other); /* 它们的形状是否能够点积 */
|
||||
auto matrix = Matrx<T>(this->get_row(), other.get_col());
|
||||
for (uint64_t i=0; i<this->get_row(); ++i) /* 点积运算 */
|
||||
for (uint64_t j=0; j<other.get_col(); ++j) {
|
||||
T sum{0};
|
||||
for (uint64_t k=0; k<this->get_col(); ++k)
|
||||
sum += this->get(i, k) * other.get(k, j);
|
||||
matrix.witch(i, j) = sum;
|
||||
}
|
||||
return matrix;
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
Matrx<T> Matrx<T>::operator*(T& item) { /* 操作符重载 - 矩阵乘法广播 */
|
||||
std::vector<T> vec = this->get_data();
|
||||
for (auto it = vec.begin(); it != vec.end(); ++it)
|
||||
*it *= item;
|
||||
return Matrx<T>{this->get_row(), this->get_col(), vec};
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
Matrx<T> Matrx<T>::operator*(const Matrx<T>& other) { /* 加法操作符重定义 - 矩阵加法 */
|
||||
//assert_same_shape(other); /* 检查两个矩阵形状是否一致 */
|
||||
std::vector<T> v0 = this->get_data();
|
||||
std::vector<T> vec = other.get_data();
|
||||
for (auto it = v0.begin(), it_or = vec.begin(); it != v0.end(); ++it, ++it_or )
|
||||
*it *= *it_or;
|
||||
return Matrx<T>{this->get_row(), this->get_col(), v0};
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
Matrx<T> Matrx<T>::operator/(T& item) { /* 操作符重载 - 矩阵除法广播 */
|
||||
std::vector<T> vec = this->get_data();
|
||||
for (auto it = vec.begin(); it != vec.end(); ++it)
|
||||
*it /= item;
|
||||
return Matrx<T>(this->get_row(), this->get_col(), vec);
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
Matrx<T> Matrx<T>::operator/(const Matrx<T>& other) { /* 加法操作符重定义 - 矩阵加法 */
|
||||
//assert_same_shape(other); /* 检查两个矩阵形状是否一致 */
|
||||
std::vector<T> v0 = this->get_data();
|
||||
std::vector<T> vec = other.get_data();
|
||||
for (auto it = v0.begin(), it_or = vec.begin(); it != v0.end(); ++it, ++it_or )
|
||||
*it /= *it_or;
|
||||
return Matrx<T>{this->get_row(), this->get_col(), v0};
|
||||
}
|
||||
|
||||
template class Matrx<float>; /* 显式实例化 */
|
||||
template class Matrx<double>;
|
||||
|
||||
}
|
||||
|
||||
/* FIXME: 形状不合法时返回空指针 */
|
||||
using namespace qing;
|
||||
extern "C" {
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void* createMatrix(uint64_t row, uint64_t col) {
|
||||
return (void *) new Matrx<float>(row, col);
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void freeMatrix(void* M) {
|
||||
delete (Matrx<float>*) M;
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
float* witchMatrix(void *m, uint64_t row, uint64_t col) {
|
||||
Matrx<float> *M = (Matrx<float>*) m;
|
||||
return &(M->witch(row, col));
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
float getMatrix(void *m, uint64_t row, uint64_t col) {
|
||||
Matrx<float> *M = (Matrx<float>*) m;
|
||||
return M->get(row, col);
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void setMatrix(void *m, uint64_t row, uint64_t col, float v) {
|
||||
Matrx<float> *M = (Matrx<float>*) m;
|
||||
M->set(row, col, v);
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
uint64_t getrowMatrix(void *m) {
|
||||
Matrx<float> *M = (Matrx<float>*) m;
|
||||
return M->get_row();
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
uint64_t getcolMatrix(void *m) {
|
||||
Matrx<float> *M = (Matrx<float>*) m;
|
||||
return M->get_col();
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
bool sameShapeMatrix(void *m, void *s) {
|
||||
Matrx<float> *M = (Matrx<float>*) m;
|
||||
Matrx<float> *S = (Matrx<float>*) s;
|
||||
return M->has_same_shape(*S);
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
bool dotShapeMatrix(void *m, void *s) {
|
||||
Matrx<float> *M = (Matrx<float>*) m;
|
||||
Matrx<float> *S = (Matrx<float>*) s;
|
||||
return M->has_dot_shape(*S);
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void* plusMatrix(void *m, void *s) {
|
||||
Matrx<float> *M = (Matrx<float>*) m;
|
||||
Matrx<float> *S = (Matrx<float>*) s;
|
||||
auto a = new Matrx<float>((*M) + (*S));
|
||||
return (void *) a;
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void* plusBroadcastMatrix(void *m, float t) {
|
||||
Matrx<float> *M = (Matrx<float>*) m;
|
||||
auto a = new Matrx<float> ((*M) + t);
|
||||
return (void*) a;
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void* subMatrix(void *m, void *s) {
|
||||
Matrx<float> *M = (Matrx<float>*) m;
|
||||
Matrx<float> *S = (Matrx<float>*) s;
|
||||
auto a = new Matrx<float>((*M) - (*S));
|
||||
return (void *) a;
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void * subBroadcastMatrix(void *m, float t) {
|
||||
Matrx<float> *M = (Matrx<float> *) m;
|
||||
auto a = new Matrx<float> ((*M) - t);
|
||||
return (void *) a;
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void* multiMatrix(void *m, void *s) {
|
||||
Matrx<float> *M = (Matrx<float>*) m;
|
||||
Matrx<float> *S = (Matrx<float>*) s;
|
||||
auto a = new Matrx<float>((*M) * (*S));
|
||||
return (void *) a;
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void* multiBroadcastMatrix(void *m, float t) {
|
||||
Matrx<float> *M = (Matrx<float> *) m;
|
||||
auto a = new Matrx<float>((*M) * t);
|
||||
return (void*) a;
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void* divMatrix(void *m, void *s) {
|
||||
Matrx<float> *M = (Matrx<float>*) m;
|
||||
Matrx<float> *S = (Matrx<float>*) s;
|
||||
auto a = new Matrx<float>((*M) / (*S));
|
||||
return (void *) a;
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void* divBroadcastMatrix(void *m, float t) {
|
||||
Matrx<float> *M = (Matrx<float> *) m;
|
||||
auto a = new Matrx<float>((*M) / t);
|
||||
return (void*) a;
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void* dotMatrix(void *m, void *s) {
|
||||
Matrx<float> *M = (Matrx<float>*) m;
|
||||
Matrx<float> *S = (Matrx<float>*) s;
|
||||
auto a = new Matrx<float>(M->dot(*S));
|
||||
return (void *) a;
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void* transposeMatrix(void *m) {
|
||||
Matrx<float> *M = (Matrx<float> *)m;
|
||||
auto N= new Matrx<float>(M->get_col(), M->get_row());
|
||||
for (uint64_t row = 0; row< N->get_row(); row++) {
|
||||
for (uint64_t col = 0; col < N->get_col(); col++) {
|
||||
N->witch(row, col) = M->get(col, row);
|
||||
}
|
||||
}
|
||||
return (void *)N;
|
||||
}
|
||||
|
||||
/* 正态分布Xavier初始化 */
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void* xavierNormalInit(uint64_t fan_in, uint64_t fan_out)
|
||||
{
|
||||
/* 构建空数组 */
|
||||
std::vector<float> weights(fan_in * fan_out);
|
||||
|
||||
/* 高质量的随机数生成器 */
|
||||
static std::mt19937 generator(std::random_device{}());
|
||||
|
||||
/* 计算标准差 */
|
||||
float sigma = std::sqrt(2.0f / (fan_in + fan_out));
|
||||
|
||||
/* 创建正态分布对象,均值为0,标准差为sigma */
|
||||
std::normal_distribution<float> distribution(0.0f, sigma);
|
||||
|
||||
/* 遍历权重矩阵,进行赋值 */
|
||||
uint64_t k=0;
|
||||
for (uint64_t i=0; i<fan_in; ++i)
|
||||
for (uint64_t j=0; j<fan_out; ++j)
|
||||
weights[k++] = distribution(generator);
|
||||
|
||||
/* 构建矩阵 */
|
||||
auto matrix = new Matrx<float>(fan_in, fan_out, weights);
|
||||
return (void *)matrix;
|
||||
}
|
||||
|
||||
/* 均匀分布Xavier初始化函数 */
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void* xavierUniformInit(uint64_t fan_in, uint64_t fan_out)
|
||||
{
|
||||
/* 构建空的数组 */
|
||||
auto weights = std::vector<float>(fan_in * fan_out);
|
||||
|
||||
/* 使用高质量的随机数生成器 */
|
||||
static std::mt19937 generator(std::random_device{}());
|
||||
|
||||
/* 计算区间范围:sqrt(2 / (fan_in + fan_out)) */
|
||||
float scale = std::sqrt(2.0f / (fan_in + fan_out));
|
||||
|
||||
/* 创建均匀分布对象,范围为[-scale, scale] */
|
||||
std::uniform_real_distribution<float> distribution(-scale, scale);
|
||||
|
||||
/* 遍历权重矩阵,进行赋值 */
|
||||
uint64_t k=0;
|
||||
for (uint64_t i=0; i<fan_in; ++i)
|
||||
for (uint64_t j=0; j<fan_out; ++j)
|
||||
weights[k++] = distribution(generator);
|
||||
|
||||
/* 构建矩阵对象 */
|
||||
auto matrix = new Matrx<float> (fan_in, fan_out, weights);
|
||||
return (void *)matrix;
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void* Sigmoid_do(void *m) {
|
||||
Matrx<float> *M = (Matrx<float>*) m;
|
||||
std::vector<float> v = M->get_data();
|
||||
for (float& x: v) {
|
||||
x = 1.0 / (1.0 + std::exp(-x));
|
||||
}
|
||||
auto N = new Matrx<float> (M->get_row(), M->get_col(), v);
|
||||
return (void *)N;
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void* Sigmoid_dot(void *m) {
|
||||
Matrx<float> *M = (Matrx<float>*) m;
|
||||
std::vector<float> v = M->get_data();
|
||||
for (float& o: v) {
|
||||
o = o * (1.0 - o);
|
||||
}
|
||||
auto N = new Matrx<float> (M->get_row(), M->get_col(), v);
|
||||
return (void *)N;
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void* ReLU_do(void *m) {
|
||||
Matrx<float> *M = (Matrx<float>*) m;
|
||||
std::vector<float> v = M->get_data();
|
||||
for (float& x: v) {
|
||||
x = x > 0 ? x : 0;
|
||||
}
|
||||
auto N = new Matrx<float> (M->get_row(), M->get_col(), v);
|
||||
return (void *)N;
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void* ReLU_dot(void *m) {
|
||||
Matrx<float> *M = (Matrx<float>*) m;
|
||||
std::vector<float> v = M->get_data();
|
||||
for (float& x: v) {
|
||||
x = x > 0 ? 1.0f : 0.0f;
|
||||
}
|
||||
auto N = new Matrx<float> (M->get_row(), M->get_col(), v);
|
||||
return (void *)N;
|
||||
}
|
||||
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void* Leaky_ReLU_do(void *m, float alpha) {
|
||||
Matrx<float> *M = (Matrx<float>*) m;
|
||||
std::vector<float> v = M->get_data();
|
||||
for (float& x: v) {
|
||||
x = x > 0 ? x : alpha * x;
|
||||
}
|
||||
auto N = new Matrx<float> (M->get_row(), M->get_col(), v);
|
||||
return (void *)N;
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void* Leaky_ReLU_dot(void *m, float alpha) {
|
||||
Matrx<float> *M = (Matrx<float>*) m;
|
||||
std::vector<float> v = M->get_data();
|
||||
for (float& x: v) {
|
||||
x = x > 0 ? 1.0 : alpha;
|
||||
}
|
||||
auto N = new Matrx<float> (M->get_row(), M->get_col(), v);
|
||||
return (void *)N;
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void* Tanh_do(void *m) {
|
||||
Matrx<float> *M = (Matrx<float>*) m;
|
||||
std::vector<float> v = M->get_data();
|
||||
for (float& x: v) {
|
||||
x = std::tanh(x);
|
||||
}
|
||||
auto N = new Matrx<float> (M->get_row(), M->get_col(), v);
|
||||
return (void *)N;
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
_declspec(dllexport)
|
||||
#endif
|
||||
void* Tanh_dot(void *m) {
|
||||
Matrx<float> *M = (Matrx<float>*) m;
|
||||
std::vector<float> v = M->get_data();
|
||||
for (float& x: v) {
|
||||
if (x > 20.0) x = 0.0; // tanh(x) ≈ 1, 导数≈0
|
||||
else if (x < -20.0) x = 0.0; // tanh(x) ≈ -1, 导数≈0
|
||||
else {
|
||||
float t = tanh(x);
|
||||
x = 1.0 - t * t;
|
||||
}
|
||||
}
|
||||
auto N = new Matrx<float> (M->get_row(), M->get_col(), v);
|
||||
return (void *)N;
|
||||
}
|
||||
}
|
||||
@@ -1,60 +0,0 @@
|
||||
/*
|
||||
|
||||
___col___
|
||||
/ \
|
||||
+---------+ \
|
||||
| + + + + | | row - 第一维度
|
||||
| + + + + | |
|
||||
+---------+ /
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
#include <cmath>
|
||||
#include <vector>
|
||||
#include <iostream>
|
||||
#include <stdexcept>
|
||||
#include <cstdint>
|
||||
namespace qing {
|
||||
template <class T>
|
||||
class Matrx {
|
||||
public:
|
||||
Matrx();
|
||||
Matrx(const uint64_t row, const uint64_t col);
|
||||
Matrx(const uint64_t row, const uint64_t col, const std::vector<T>& src);
|
||||
const T& get(uint64_t row, uint64_t col) const;
|
||||
//-------------旧接口----------------
|
||||
T& witch(uint64_t row, uint64_t col);
|
||||
//-------------新接口----------------
|
||||
void set(uint64_t row, uint64_t col, T v);
|
||||
//-----------------------------------
|
||||
const std::vector<T>& get_data() const;
|
||||
const uint64_t get_col() const;
|
||||
const uint64_t get_row() const;
|
||||
bool has_same_shape(const Matrx<T>& other);
|
||||
bool has_dot_shape(const Matrx<T>& other);
|
||||
void assert_same_shape(const Matrx<T>& other);
|
||||
void assert_dot_shape(const Matrx<T>& other);
|
||||
Matrx<T> operator+(const Matrx<T>& other);
|
||||
Matrx<T> operator+(T& item);
|
||||
Matrx<T> operator-(const Matrx<T>& other);
|
||||
Matrx<T> operator-(T& item);
|
||||
Matrx<T> operator*(const Matrx<T>& other);
|
||||
Matrx<T> operator*(T& item);
|
||||
Matrx<T> operator/(T& item);
|
||||
Matrx<T> operator/(const Matrx<T>& other);
|
||||
Matrx<T> dot(const Matrx<T>& other);
|
||||
|
||||
void save(std::ostream& out) {
|
||||
out << row << " " << col;
|
||||
for (const auto& item : arr) {
|
||||
out << " ";
|
||||
out << item;
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
uint64_t row, col;
|
||||
std::vector<T> arr;
|
||||
};
|
||||
}
|
||||
@@ -1,21 +0,0 @@
|
||||
/* I2C操作库 - PCA9685 - MPU6050 */
|
||||
#ifndef QING_I2C
|
||||
#define QING_I2C
|
||||
/* I2C */
|
||||
#define BOT4CJ_I2C_OK 0
|
||||
#define BOT4CJ_I2C_ERR 1
|
||||
#define BOT4CJ_I2C_OPENED_FAILED -1
|
||||
|
||||
#ifdef _WIN64
|
||||
#define DLLEXPORT __declspec(dllexport)
|
||||
#else
|
||||
#define DLLEXPORT
|
||||
#endif
|
||||
|
||||
DLLEXPORT int16_t I2C_init(const char* filepath);
|
||||
DLLEXPORT void I2C_free(int16_t fd);
|
||||
DLLEXPORT int I2C_setAddr(int16_t fd, uint8_t addr);
|
||||
DLLEXPORT int I2C_writeReg(int16_t fd, uint8_t reg, uint8_t val);
|
||||
DLLEXPORT uint8_t I2C_readReg(int16_t fd, uint8_t reg);
|
||||
DLLEXPORT int8_t I2C_readRegFlag(int16_t fd, uint8_t reg);
|
||||
#endif
|
||||
+37
-62
@@ -1,15 +1,20 @@
|
||||
package bot4cj
|
||||
|
||||
/*
|
||||
* 机器人脚本解析器 - 仓颉语言
|
||||
*
|
||||
*/
|
||||
|
||||
import std.argopt.*
|
||||
import std.collection.*
|
||||
import std.env.*
|
||||
import std.convert.*
|
||||
import std.io.*
|
||||
import std.fs.*
|
||||
import std.env.*
|
||||
import std.time.*
|
||||
import std.argopt.*
|
||||
import std.convert.*
|
||||
import std.collection.*
|
||||
import stdx.encoding.json.*
|
||||
import bot4cj.hd.*
|
||||
import bot4cj.math.*
|
||||
import ctype4cj.*
|
||||
import nn4cj.*
|
||||
|
||||
/* ---- 神经网络输出路径 ---- */
|
||||
var NNS_OUT: Option<String> = None
|
||||
@@ -17,10 +22,10 @@ main(args: Array<String>): Int64 {
|
||||
|
||||
/* ---- 解析命令行参数 ---- */
|
||||
let specs = [
|
||||
ArgumentSpec.Long("help", ArgumentMode.NoValue), /* ---- 拉起帮助菜单 ---- */
|
||||
ArgumentSpec.Short(r'i', ArgumentMode.RequiredValue), /* ---- 神经网络模型 - 输入 ---- */
|
||||
ArgumentSpec.Short(r'o', ArgumentMode.RequiredValue), /* ---- 神经网络模型 - 输出 ---- */
|
||||
ArgumentSpec.Long("test", ArgumentMode.RequiredValue) /* ---- 运行一个测试分支 - 需要输入测试类型 ---- */
|
||||
ArgumentSpec.Long("help", ArgumentMode.NoValue), /* ---- 拉起帮助 ---- */
|
||||
ArgumentSpec.Short(r'i', ArgumentMode.RequiredValue), /* ---- 输入模型 ---- */
|
||||
ArgumentSpec.Short(r'o', ArgumentMode.RequiredValue), /* ---- 输出模型 ---- */
|
||||
ArgumentSpec.Long("test", ArgumentMode.RequiredValue) /* ---- 输入一个测试 ---- */
|
||||
]
|
||||
|
||||
/* ---- 执行解析 ---- */
|
||||
@@ -30,10 +35,10 @@ main(args: Array<String>): Int64 {
|
||||
//---------------------------
|
||||
/* 读取设置项 */
|
||||
|
||||
/*---------------------------*/
|
||||
/*---------------------------*/
|
||||
/* 初始化神经网络 (init nn)
|
||||
* TODO: 可以去掉内部初始化参数,使用矩阵式模板进行初始化 */
|
||||
/*
|
||||
* 初始化神经网络 (init nn)
|
||||
* TODO: 可以去掉内部初始化参数,使用矩阵式模板进行初始化
|
||||
*/
|
||||
|
||||
|
||||
//var nns: ArrayList<Nn> =
|
||||
@@ -53,28 +58,27 @@ main(args: Array<String>): Int64 {
|
||||
NNS_OUT = Option<String>.Some(options['o'])
|
||||
}
|
||||
|
||||
/*---------------------------*/
|
||||
/*---------------------------*/
|
||||
/* 分支选择 */
|
||||
|
||||
|
||||
/* -------- 拉起帮助手册
|
||||
* -------- */
|
||||
/* -------- 拉起帮助手册 -------- */
|
||||
if (options.contains("help")) {
|
||||
print_help()
|
||||
return 0
|
||||
}
|
||||
|
||||
/* -------- 启动一个测试分支
|
||||
* -------- */
|
||||
/* -------- 启动一个测试分支 -------- */
|
||||
else if (options.contains("test")) {
|
||||
run_test(options["test"])
|
||||
return 0
|
||||
}
|
||||
|
||||
/* -------- 缺省行为:从标准输入获取动作脚本
|
||||
* --------
|
||||
* -------- TODO(20260519-qing): 进行脚本引入(相对路径也可,从**环境变量**的路径中搜寻也可) */
|
||||
/*
|
||||
* 缺省行为:从标准输入获取动作脚本
|
||||
*
|
||||
* TODO(20260519-qing): 进行脚本引入(相对路径也可,从**环境变量**的路径中搜寻也可)
|
||||
*/
|
||||
else {
|
||||
let consoleReader = getStdIn()
|
||||
let robot = RobotStruct(consoleReader)
|
||||
@@ -82,12 +86,11 @@ main(args: Array<String>): Int64 {
|
||||
robot.active()
|
||||
robot.run()
|
||||
}
|
||||
finally {
|
||||
robot.deactive()
|
||||
}
|
||||
|
||||
/* ---- NOTE(20260523-qing): 储存微调网络模型
|
||||
* ---- */
|
||||
/* 无论有没有成功激活,都进行清理 */
|
||||
finally { robot.deactive() }
|
||||
|
||||
/* NOTE(20260523-qing): 储存微调网络模型 */
|
||||
// let consoleReader = getStdIn()
|
||||
// run_action_drl(consoleReader, nns)
|
||||
// if (options.contains("o")) {
|
||||
@@ -98,8 +101,7 @@ main(args: Array<String>): Int64 {
|
||||
return 0
|
||||
}
|
||||
|
||||
/* -------- 打印帮助手册
|
||||
* -------- */
|
||||
/* -------- 打印帮助手册 -------- */
|
||||
func print_help()
|
||||
{
|
||||
println("这是一个机器人动作脚本解析器。")
|
||||
@@ -112,43 +114,16 @@ func print_help()
|
||||
println("-o 设置输出的模型的路径")
|
||||
}
|
||||
|
||||
/* -------- 集成多种常规测试、临时测试
|
||||
* -------- */
|
||||
/* -------- 集成多种常规测试、临时测试 -------- */
|
||||
func run_test(key: String) {
|
||||
|
||||
/* ----------- 神经网络与门测试
|
||||
* -----------
|
||||
* ----------- INFO: 我们使用随机生成的动态数据集,这样会导致一些后果 */
|
||||
if (key == "nn") {
|
||||
let nn_t = ArrayList<Nn1>()
|
||||
nn_t.add(Nn1(2, 4, 0.05f32, Af.Leaky_ReLU(FP32(0.01))))
|
||||
nn_t.add(Nn1(4, 1, 0.05f32, Af.Sigmoid))
|
||||
let batch = 1
|
||||
let num =2000
|
||||
let prt_num = 100
|
||||
for (i in 0..(num*batch)) {
|
||||
let m = Matrix1(1, 2)
|
||||
let t = Matrix1(1, 1)
|
||||
prepare_andDoor(m, t)
|
||||
let l = train1(nn_t, m, t, 1.0f32) // 这个1.0是学习率的折扣,而非学习率本身
|
||||
if (i % prt_num == prt_num-1) {
|
||||
print("${i}\t:")
|
||||
print("${l}")
|
||||
println()
|
||||
}
|
||||
}
|
||||
/*
|
||||
* 神经网络门测试
|
||||
*/
|
||||
if (key == "nn1") {
|
||||
testNn()
|
||||
}
|
||||
|
||||
|
||||
else if (key == "cnn") {
|
||||
testCNN()
|
||||
}
|
||||
|
||||
else if (key == "cnn1") {
|
||||
testGeometryCNN()
|
||||
}
|
||||
|
||||
|
||||
/* -------- 测试解析矩阵,并且输出
|
||||
* -------- */
|
||||
else if (key == "matrix_o") {
|
||||
|
||||
@@ -1,204 +0,0 @@
|
||||
// Adam.cj - Adam 优化器实现(添加全连接层支持)
|
||||
package bot4cj.math
|
||||
|
||||
import std.math.*
|
||||
import std.collection.*
|
||||
|
||||
public class Adam {
|
||||
let learningRate: FP32
|
||||
let beta1: FP32
|
||||
let beta2: FP32
|
||||
let epsilon: FP32
|
||||
|
||||
var m: HashMap<String, Matrix1>
|
||||
var v: HashMap<String, Matrix1>
|
||||
var t: Int64
|
||||
|
||||
// 完整参数构造函数
|
||||
public init(
|
||||
lr: FP32,
|
||||
beta1: FP32,
|
||||
beta2: FP32,
|
||||
epsilon: FP32
|
||||
) {
|
||||
learningRate = lr
|
||||
this.beta1 = beta1
|
||||
this.beta2 = beta2
|
||||
this.epsilon = epsilon
|
||||
m = HashMap<String, Matrix1>()
|
||||
v = HashMap<String, Matrix1>()
|
||||
t = 0
|
||||
}
|
||||
|
||||
// 默认参数构造函数
|
||||
public init() {
|
||||
learningRate = FP32(0.001)
|
||||
beta1 = FP32(0.9)
|
||||
beta2 = FP32(0.999)
|
||||
epsilon = FP32(0.00000001)
|
||||
m = HashMap<String, Matrix1>()
|
||||
v = HashMap<String, Matrix1>()
|
||||
t = 0
|
||||
}
|
||||
|
||||
// 仅自定义学习率的构造函数
|
||||
public init(lr: FP32) {
|
||||
learningRate = lr
|
||||
beta1 = FP32(0.9)
|
||||
beta2 = FP32(0.999)
|
||||
epsilon = FP32(0.00000001)
|
||||
m = HashMap<String, Matrix1>()
|
||||
v = HashMap<String, Matrix1>()
|
||||
t = 0
|
||||
}
|
||||
|
||||
/**
|
||||
* 计算幂运算的辅助函数
|
||||
*/
|
||||
private func power(base: FP32, exp: Int64): FP32 {
|
||||
var result = FP32(1.0)
|
||||
for (i in 0..exp) {
|
||||
result = result * base
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
/**
|
||||
* 计算平方根的辅助函数 - 使用牛顿迭代法
|
||||
*/
|
||||
private func sqrt(val: FP32): FP32 {
|
||||
let x = val.getValue()
|
||||
if (x <= 0.0) {
|
||||
return FP32(0.0)
|
||||
}
|
||||
var guess = x
|
||||
for (i in 0..10) {
|
||||
guess = (guess + x / guess) / 2.0
|
||||
}
|
||||
return FP32(guess)
|
||||
}
|
||||
|
||||
/**
|
||||
* 更新矩阵参数(通用方法)
|
||||
* @param param 要更新的参数矩阵
|
||||
* @param grad 梯度矩阵
|
||||
* @param key 参数在HashMap中的键名
|
||||
*/
|
||||
private func updateMatrix(
|
||||
param: Matrix1,
|
||||
grad: Matrix1,
|
||||
key: String
|
||||
): Unit {
|
||||
// 初始化动量
|
||||
let hasState = m.contains(key)
|
||||
if (!hasState) {
|
||||
let rows = param.get_row()
|
||||
let cols = param.get_col()
|
||||
m.add(key, Matrix1(rows, cols))
|
||||
v.add(key, Matrix1(rows, cols))
|
||||
}
|
||||
|
||||
let mMat = m.get(key).getOrThrow()
|
||||
let vMat = v.get(key).getOrThrow()
|
||||
|
||||
let rows = param.get_row()
|
||||
let cols = param.get_col()
|
||||
|
||||
for (r in 0..rows) {
|
||||
for (c in 0..cols) {
|
||||
let g = grad.get(r, c).getOrThrow()
|
||||
|
||||
// 更新一阶矩
|
||||
let oldM = mMat.get(r, c).getOrThrow()
|
||||
let newM = beta1 * oldM + (FP32(1.0) - beta1) * g
|
||||
mMat.set(r, c, newM)
|
||||
|
||||
// 更新二阶矩
|
||||
let oldV = vMat.get(r, c).getOrThrow()
|
||||
let newV = beta2 * oldV + (FP32(1.0) - beta2) * g * g
|
||||
vMat.set(r, c, newV)
|
||||
|
||||
// 偏差修正
|
||||
let one = FP32(1.0)
|
||||
let beta1Pow = power(beta1, t)
|
||||
let beta2Pow = power(beta2, t)
|
||||
|
||||
let mHat = newM / (one - beta1Pow)
|
||||
let vHat = newV / (one - beta2Pow)
|
||||
|
||||
// 更新参数
|
||||
let oldParam = param.get(r, c).getOrThrow()
|
||||
let update = learningRate * mHat / (sqrt(vHat) + epsilon)
|
||||
param.set(r, c, oldParam - update)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 更新卷积层参数
|
||||
*/
|
||||
public func updateConvLayer(layer: ConvLayer): Unit {
|
||||
t += 1
|
||||
|
||||
// 更新卷积核
|
||||
var kGrad = match (layer.kernelsGrad) {
|
||||
case Some(val) => val
|
||||
case _ => throw Exception("No kernel gradients")
|
||||
}
|
||||
|
||||
for (oc in 0..layer.nOutputChannels) {
|
||||
let key = "conv_kernel_" + oc.toString()
|
||||
updateMatrix(layer.kernels[oc], kGrad[oc], key)
|
||||
}
|
||||
|
||||
// 更新偏置
|
||||
var bGrad = match (layer.biasGrad) {
|
||||
case Some(val) => val
|
||||
case _ => throw Exception("No bias gradient")
|
||||
}
|
||||
|
||||
let biasKey = "conv_bias"
|
||||
updateMatrix(layer.biases, bGrad, biasKey)
|
||||
}
|
||||
|
||||
/**
|
||||
* 更新全连接层参数(Nn1)
|
||||
*/
|
||||
public func updateFCLayer(layer: Nn1): Unit {
|
||||
t += 1
|
||||
|
||||
// 更新权重
|
||||
var wGrad = match (layer.weights_grad) {
|
||||
case Some(val) => val
|
||||
case _ => throw Exception("No weights gradient")
|
||||
}
|
||||
|
||||
let weightKey = "fc_weights"
|
||||
updateMatrix(layer.weights, wGrad, weightKey)
|
||||
|
||||
// 更新偏置
|
||||
var bGrad = match (layer.bias_grad) {
|
||||
case Some(val) => val
|
||||
case _ => throw Exception("No bias gradient")
|
||||
}
|
||||
|
||||
let biasKey = "fc_bias"
|
||||
updateMatrix(layer.bias, bGrad, biasKey)
|
||||
}
|
||||
|
||||
/**
|
||||
* 重置优化器状态
|
||||
*/
|
||||
public func reset(): Unit {
|
||||
m.clear()
|
||||
v.clear()
|
||||
t = 0
|
||||
}
|
||||
|
||||
/**
|
||||
* 获取当前迭代次数
|
||||
*/
|
||||
public func getStep(): Int64 {
|
||||
t
|
||||
}
|
||||
}
|
||||
@@ -1,273 +0,0 @@
|
||||
package bot4cj.math
|
||||
import std.math.*
|
||||
|
||||
/* -------- Activation Functions 激活函数族 --------
|
||||
* TODO(20260516-qing): 增加序列化、反序列化工厂函数 */
|
||||
public enum Af {
|
||||
|
||||
None | Sigmoid | ReLU | Leaky_ReLU(FP32) | Tanh | Softmax
|
||||
|
||||
/* -------- 执行激活函数 -------- */
|
||||
public func forward(x: Matrix1) {
|
||||
match(this) {
|
||||
case None => none(x)
|
||||
case Sigmoid => sigmoid(x)
|
||||
case ReLU => relu(x)
|
||||
case Leaky_ReLU(alpha) => leaky_relu(x, alpha)
|
||||
case Tanh => tanh(x)
|
||||
case Softmax => softmax(x)
|
||||
}
|
||||
}
|
||||
|
||||
/* -------- 激活函数的导数 -------- */
|
||||
public func backward(y: Matrix1) {
|
||||
match(this) {
|
||||
case None => none_dot(y)
|
||||
case Sigmoid => sigmoid_dot(y)
|
||||
case ReLU => relu_dot(y)
|
||||
case Leaky_ReLU(alpha) => leaky_relu_dot(y, alpha)
|
||||
case Tanh => tanh_dot(y)
|
||||
case Softmax => softmax_jacobian(y)
|
||||
}
|
||||
}
|
||||
|
||||
/* -------- 序列化与反序列化 -------- */
|
||||
public func get_type(): String {
|
||||
return match(this) {
|
||||
case None => "None"
|
||||
case Sigmoid => "Sigmoid"
|
||||
case ReLU => "ReLU"
|
||||
case Leaky_ReLU(alpha) => "Leaky_ReLU"
|
||||
case Tanh => "Tanh"
|
||||
case Softmax => "Softmax"
|
||||
}
|
||||
}
|
||||
|
||||
/* */
|
||||
public func get_value(): FP32 {
|
||||
return match(this) {
|
||||
case None => FP32(0.0f32)
|
||||
case Sigmoid => FP32(0.0f32)
|
||||
case ReLU => FP32(0.0f32)
|
||||
case Leaky_ReLU(alpha) => alpha
|
||||
case Tanh => FP32(0.0f32)
|
||||
case Softmax => FP32(0.0f32)
|
||||
}
|
||||
}
|
||||
|
||||
/* -------- 空的激活函数 -------- */
|
||||
private static func none(M: Matrix1) {
|
||||
let row = M.get_row()
|
||||
let col = M.get_col()
|
||||
let N = Matrix1(row, col)
|
||||
for (r in 0..row) {
|
||||
for (c in 0..col) {
|
||||
let x = M.get(r, c).getOrThrow()
|
||||
let x1= x.getValue()
|
||||
N.set(r, c, FP32(x1)) }
|
||||
}
|
||||
return N
|
||||
}
|
||||
|
||||
/* -------- 空激活函数的导数 -------- */
|
||||
private static func none_dot(M: Matrix1) {
|
||||
let row = M.get_row()
|
||||
let col = M.get_col()
|
||||
let N = Matrix1(row, col)
|
||||
for (r in 0..row) {
|
||||
for (c in 0..col) {
|
||||
N.set(r, c, FP32(1.0f32)) }
|
||||
}
|
||||
return N
|
||||
}
|
||||
|
||||
/* -------- S型生长曲线 -------- */
|
||||
private static func sigmoid(M: Matrix1) {
|
||||
let row = M.get_row()
|
||||
let col = M.get_col()
|
||||
let N = Matrix1(row, col)
|
||||
for (r in 0..row) {
|
||||
for (c in 0..col) {
|
||||
let x = M.get(r, c).getOrThrow()
|
||||
let x1= x.getValue()
|
||||
let x2 = 1.0f32 / (1.0f32 + exp(-x1))
|
||||
N.set(r, c, FP32(x2)) }
|
||||
}
|
||||
return N
|
||||
}
|
||||
|
||||
/* -------- S型生长曲线的导数 -------- */
|
||||
private static func sigmoid_dot(M: Matrix1) {
|
||||
let row = M.get_row()
|
||||
let col = M.get_col()
|
||||
let N = Matrix1(row, col)
|
||||
for (r in 0..row) {
|
||||
for (c in 0..col) {
|
||||
let x = M.get(r, c).getOrThrow()
|
||||
let x1= x.getValue()
|
||||
let x2 = 1.0f32 / (1.0f32 + exp(-x1))
|
||||
let x3 = x2 * (1.0f32 - x2)
|
||||
N.set(r, c, FP32(x3)) }
|
||||
}
|
||||
return N
|
||||
}
|
||||
|
||||
/* -------- 线性整流单元 -------- */
|
||||
private static func relu(M: Matrix1) {
|
||||
let row = M.get_row()
|
||||
let col = M.get_col()
|
||||
let N = Matrix1(row, col)
|
||||
for (r in 0..row) {
|
||||
for (c in 0..col) {
|
||||
let x = M.get(r, c).getOrThrow()
|
||||
let x1= x.getValue()
|
||||
let x2 =if (x1 > 0.0f32) {x1} else{0.0f32}
|
||||
N.set(r, c, FP32(x2))
|
||||
}
|
||||
}
|
||||
return N
|
||||
}
|
||||
|
||||
/* -------- 线性整流单元的导数 -------- */
|
||||
private static func relu_dot(M: Matrix1) {
|
||||
let row = M.get_row()
|
||||
let col = M.get_col()
|
||||
let N = Matrix1(row, col)
|
||||
for (r in 0..row) {
|
||||
for (c in 0..col) {
|
||||
let x = M.get(r, c).getOrThrow()
|
||||
let x1= x.getValue()
|
||||
let x2 =if (x1 > 0.0f32) {1.0f32} else{0.0f32}
|
||||
N.set(r, c, FP32(x2))
|
||||
}
|
||||
}
|
||||
return N
|
||||
}
|
||||
|
||||
/* -------- 带泄露的线性整流单元 -------- */
|
||||
private static func leaky_relu(M: Matrix1, alpha: FP32) {
|
||||
let row = M.get_row()
|
||||
let col = M.get_col()
|
||||
let N = Matrix1(row, col)
|
||||
for (r in 0..row) {
|
||||
for (c in 0..col) {
|
||||
let x = M.get(r, c).getOrThrow()
|
||||
let x1= x.getValue()
|
||||
let x2 =if (x1 > 0.0f32) {x1} else{x1 * alpha.getValue()}
|
||||
N.set(r, c, FP32(x2))
|
||||
}
|
||||
}
|
||||
return N
|
||||
}
|
||||
|
||||
/* -------- 带泄露的整流单元导数 -------- */
|
||||
private static func leaky_relu_dot(M: Matrix1, alpha: FP32) {
|
||||
let row = M.get_row()
|
||||
let col = M.get_col()
|
||||
let N = Matrix1(row, col)
|
||||
for (r in 0..row) {
|
||||
for (c in 0..col) {
|
||||
let x = M.get(r, c).getOrThrow()
|
||||
let x1= x.getValue()
|
||||
let x2 =if (x1 > 0.0f32) { 1.0f32 } else{alpha.getValue()}
|
||||
N.set(r, c, FP32(x2)) }
|
||||
}
|
||||
return N
|
||||
}
|
||||
|
||||
/* -------- 双曲正切函数 --------*/
|
||||
private static func tanh(M: Matrix1) {
|
||||
let row = M.get_row()
|
||||
let col = M.get_col()
|
||||
let N = Matrix1(row, col)
|
||||
for (r in 0..row) {
|
||||
for (c in 0..col) {
|
||||
let x = M.get(r, c).getOrThrow()
|
||||
let x1= x.getValue()
|
||||
let x2 = tanh(x1)
|
||||
N.set(r, c, FP32(x2)) }
|
||||
}
|
||||
return N
|
||||
}
|
||||
|
||||
/* -------- 双曲正切函数的导数 -------- */
|
||||
private static func tanh_dot(M: Matrix1) {
|
||||
let row = M.get_row()
|
||||
let col = M.get_col()
|
||||
let N = Matrix1(row, col)
|
||||
for (r in 0..row) {
|
||||
for (c in 0..col) {
|
||||
let y = M.get(r, c).getOrThrow()
|
||||
let y_val= y.getValue() /*---- NOTE(20260525-qing): 修复命名(x => y) ----*/
|
||||
let rea = 1.0f32 - y_val * y_val
|
||||
N.set(r, c, FP32(rea)) }
|
||||
}
|
||||
return N
|
||||
}
|
||||
|
||||
/* -------- 归一化指数函数 -------- */
|
||||
private static func softmax(M: Matrix1) {
|
||||
let row = M.get_row()
|
||||
let col = M.get_col()
|
||||
let N = Matrix1(row, col)
|
||||
var max_val = 0.0f32
|
||||
for (r in 0..row) {
|
||||
for (c in 0..col) {
|
||||
let x = M.get(r, c).getOrThrow()
|
||||
let x1= x.getValue()
|
||||
if (x1 > max_val) {
|
||||
max_val = x1 }
|
||||
}
|
||||
}
|
||||
var sum = 0.0f32
|
||||
for (r in 0..row) {
|
||||
for (c in 0..col) {
|
||||
let x = M.get(r, c).getOrThrow()
|
||||
let x1= x.getValue()
|
||||
let x2 = exp(x1 - max_val)
|
||||
N.set(r, c, FP32(x2))
|
||||
sum += x2
|
||||
}
|
||||
}
|
||||
for (r in 0..row) {
|
||||
for (c in 0..col) {
|
||||
let x = N.get(r, c).getOrThrow()
|
||||
let x1= x.getValue()
|
||||
let x2 = x1 / sum
|
||||
N.set(r, c, FP32(x2))
|
||||
}
|
||||
}
|
||||
return N
|
||||
}
|
||||
|
||||
/* -------- 计算Softmax的雅可比矩阵 --------
|
||||
* 输入:已计算好的softmax输出(概率分布)
|
||||
* 输出:雅可比矩阵 J[i][j] = ∂y_i/∂x_j */
|
||||
private func softmax_jacobian(softmax_output: Matrix1) {
|
||||
let row = softmax_output.get_row()
|
||||
let col = softmax_output.get_col()
|
||||
let n = row * col
|
||||
var jacobian = Matrix1(n, n)
|
||||
|
||||
for (i in 0..n) {
|
||||
for (j in 0..n) {
|
||||
|
||||
let row_i = i / col /* 需要将一维索引转换为二维索引来访问Matrix1 */
|
||||
let col_i = i % col
|
||||
let row_j = j / col
|
||||
let col_j = j % col
|
||||
|
||||
let val_i = softmax_output.get(row_i, col_i).getOrThrow().getValue()
|
||||
let val_j = softmax_output.get(row_j, col_j).getOrThrow().getValue()
|
||||
|
||||
if (i == j) { /* 对角线元素: y_i * (1 - y_i) */
|
||||
jacobian.set(i, j, FP32(val_i * (1.0f32 - val_i))) }
|
||||
else { /* 非对角线元素: -y_i * y_j */
|
||||
jacobian.set(i, j, FP32(-val_i * val_j)) }
|
||||
}
|
||||
}
|
||||
|
||||
return jacobian
|
||||
}
|
||||
|
||||
}
|
||||
@@ -1,37 +0,0 @@
|
||||
package bot4cj.math
|
||||
import std.random.*
|
||||
import std.time.*
|
||||
import std.math.*
|
||||
|
||||
/* -------- 神经网络使用的权重矩阵初始化类 --------
|
||||
* TODO(20260517-qing): 将该函数改为枚举类,统一一个接口,根据模式匹配执行初始化
|
||||
* - 这样在矩阵式中可以定义初始化方式 */
|
||||
public class Ini {
|
||||
|
||||
/* -------- xavier正态初始化 -------- */
|
||||
public static func xavierNormalInit(target: Matrix1, fan_in: Int64, fan_out: Int64) {
|
||||
let seed = DateTime.now().nanosecond
|
||||
let random = Random(UInt64(seed))
|
||||
let sigma = sqrt(2.0f32 / Float32(fan_in + fan_out)) /* 计算标准差 */
|
||||
for (r in 0..fan_in) {
|
||||
for (c in 0..fan_out) {
|
||||
let v = random.nextGaussianFloat32(mean: 0.0f32, sigma:sigma)
|
||||
target.set(r, c, FP32(v))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/* -------- Glorot均匀初始化 -------- */
|
||||
public static func xavierUniformInit(target: Matrix1, fan_in: Int64, fan_out: Int64) {
|
||||
let limit = sqrt(6.0f32 / Float32(fan_in + fan_out))
|
||||
let seed = DateTime.now().nanosecond
|
||||
let random = Random(UInt64(seed))
|
||||
|
||||
for (r in 0..fan_in) {
|
||||
for (c in 0..fan_out) {
|
||||
let v = (random.nextFloat32() * 2.0f32 - 1.0f32) * limit /* 均匀分布 [-limit, limit] */
|
||||
target.set(r, c, FP32(v)) }
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
@@ -1,535 +0,0 @@
|
||||
package bot4cj.math
|
||||
import std.fs.*
|
||||
import std.io.*
|
||||
import std.convert.*
|
||||
//foreign func createMatrix(row: UInt64, col: UInt64): CPointer<Unit>
|
||||
//foreign func freeMatrix(M: CPointer<Unit>): Unit
|
||||
//foreign func witchMatrix(M: CPointer<Unit>, row: UInt64, col: UInt64): CPointer<Float32>
|
||||
//foreign func getMatrix(M: CPointer<Unit>, row: UInt64, col: UInt64): Float32
|
||||
//foreign func setMatrix(M: CPointer<Unit>, row: UInt64, col: UInt64, v: Float32): Unit
|
||||
//foreign func getrowMatrix(M: CPointer<Unit>): UInt64
|
||||
//foreign func getcolMatrix(M: CPointer<Unit>): UInt64
|
||||
//foreign func sameShapeMatrix(M: CPointer<Unit>, S: CPointer<Unit>): Bool
|
||||
//foreign func dotShapeMatrix(M: CPointer<Unit>, S: CPointer<Unit>): Bool
|
||||
//foreign func plusMatrix(M: CPointer<Unit>, S:CPointer<Unit>): CPointer<Unit>
|
||||
//foreign func plusBroadcastMatrix(M: CPointer<Unit>, T: Float32): CPointer<Unit>
|
||||
//foreign func subMatrix(M: CPointer<Unit>, S: CPointer<Unit>): CPointer<Unit>
|
||||
//foreign func subBroadcastMatrix(M: CPointer<Unit>, T: Float32): CPointer<Unit>
|
||||
//foreign func multiMatrix(M: CPointer<Unit>, S: CPointer<Unit>): CPointer<Unit>
|
||||
//foreign func multiBroadcastMatrix(M: CPointer<Unit>, T: Float32): CPointer<Unit>
|
||||
//foreign func divMatrix(M: CPointer<Unit>, S: CPointer<Unit>): CPointer<Unit>
|
||||
//foreign func divBroadcastMatrix(M: CPointer<Unit>, T: Float32): CPointer<Unit>
|
||||
//foreign func dotMatrix(M: CPointer<Unit>, S: CPointer<Unit>): CPointer<Unit>
|
||||
//foreign func transposeMatrix(M: CPointer<Unit>): CPointer<Unit>
|
||||
//foreign func xavierNormalInit(fan_in: UInt64, fan_out: UInt64): CPointer<Unit>
|
||||
//foreign func xavierUniformInit(fan_in: UInt64, fan_out: UInt64): CPointer<Unit>
|
||||
|
||||
|
||||
/* -------- 32位浮点数 -------- */
|
||||
public struct FP32 {
|
||||
private var val: Float32
|
||||
public init(t: Float32) { /* 以Float32初始化对象 */
|
||||
val = t }
|
||||
public func getValue() { /* 获取值 */
|
||||
return val }
|
||||
/* -------- 运算符重载 -------- */
|
||||
//public operator func =(t: Float32) { /*-------- this can not --------*/
|
||||
// val = t }
|
||||
public operator func +(t: Float32) {
|
||||
return FP32(val + t)}
|
||||
public operator func +(S: FP32) {
|
||||
return this + S.getValue() }
|
||||
public operator func -(t: Float32) {
|
||||
return FP32(val - t) }
|
||||
public operator func -(S: FP32) {
|
||||
return this - S.getValue() }
|
||||
public operator func *(t: Float32) {
|
||||
return FP32(val * t) }
|
||||
public operator func *(S: FP32) {
|
||||
return this * S.getValue() }
|
||||
public operator func /(t: Float32) {
|
||||
return FP32(val / t) }
|
||||
public operator func /(S: FP32) {
|
||||
return this / S.getValue() }
|
||||
|
||||
}
|
||||
|
||||
|
||||
/*
|
||||
* 新版的矩阵对象,仓颉原生实现
|
||||
* TODO(20260517-qing): 设计原生的多维数组类
|
||||
*
|
||||
|
||||
___col___
|
||||
/ \
|
||||
+---------+ \
|
||||
| + + + + | | row - 第一维度
|
||||
| + + + + | |
|
||||
+---------+ /
|
||||
|
||||
|
||||
* TODO(20260523-qing): 以某种方式统一矩阵式的字符串矩阵和这里的数值矩阵
|
||||
*/
|
||||
public class Matrix1 {
|
||||
|
||||
private var row: Int64 /* ---- 行数 ---- */
|
||||
private var col: Int64 /* ---- 列数 ---- */
|
||||
private var arr: Array<FP32> /* ---- 数据 ---- */
|
||||
|
||||
/* -------- 构造器,默认值为0 -------- */
|
||||
public Matrix1(row: Int64, col: Int64) {
|
||||
this.row = row
|
||||
this.col = col
|
||||
this.arr = Array<FP32>(row * col, {_ => FP32(0.0f32)})
|
||||
}
|
||||
|
||||
/* --------- 这几个方法是访问器 -------- */
|
||||
/* ------------------------------------- */
|
||||
/* -------- TODO(20260524-qing): 使用prop成员 -------- */
|
||||
|
||||
public func get(r: Int64, c: Int64) {
|
||||
return arr.get(r * this.col + c)
|
||||
}
|
||||
|
||||
public func set(r: Int64, c: Int64, v: FP32) {
|
||||
arr[r * this.col + c] = v
|
||||
}
|
||||
|
||||
public func get_row() { return row }
|
||||
|
||||
public func get_col() { return col }
|
||||
|
||||
/* -------- 这两个函数是比较器 -------- */
|
||||
/* ------------------------------------ */
|
||||
public func has_same_shape(other: Matrix1) {
|
||||
let a = this.get_row() == other.get_row()
|
||||
let b = this.get_col() == other.get_col()
|
||||
return a && b
|
||||
}
|
||||
|
||||
public func has_dot_shape(other: Matrix1) {
|
||||
return this.get_col() == other.get_row()
|
||||
}
|
||||
|
||||
/* -------- 这些是运算符重载 -------- */
|
||||
/* ---------------------------------- */
|
||||
public operator func +(a: FP32) { /* ---- 矩阵广播加法 ---- */
|
||||
let r =get_row()
|
||||
let c = get_col()
|
||||
let M = Matrix1(r, c)
|
||||
for (ir in 0..r) {
|
||||
for (ic in 0..c) {
|
||||
let v = this.get(ir, ic).getOrThrow()
|
||||
M.set(ir, ic, v + a)
|
||||
}
|
||||
}
|
||||
|
||||
return M
|
||||
}
|
||||
|
||||
public operator func +(S: Matrix1) { /* ---- 矩阵加法 ---- */
|
||||
if (!has_same_shape(S)) {
|
||||
throw IllegalArgumentException("Matrixes doesn't match")
|
||||
}
|
||||
let r =get_row()
|
||||
let c = get_col()
|
||||
let M = Matrix1(r, c)
|
||||
for (ir in 0..r) {
|
||||
for (ic in 0..c) {
|
||||
let v = this.get(ir, ic).getOrThrow()
|
||||
let vs = S.get(ir, ic).getOrThrow()
|
||||
M.set(ir, ic, v + vs)
|
||||
}
|
||||
}
|
||||
|
||||
return M
|
||||
}
|
||||
|
||||
public operator func -(a: FP32) { /* ---- 矩阵广播减法 ----*/
|
||||
let r = get_row()
|
||||
let c = get_col()
|
||||
let M = Matrix1(r, c)
|
||||
for (ir in 0..r) {
|
||||
for (ic in 0..c){
|
||||
let v = this.get(ir, ic).getOrThrow()
|
||||
M.set(ir, ic, v - a)
|
||||
}
|
||||
}
|
||||
return M
|
||||
}
|
||||
|
||||
public operator func -(S: Matrix1) { /* ---- 矩阵减法 ---- */
|
||||
if (!has_same_shape(S)) {
|
||||
throw IllegalArgumentException("Matrixes doesn't match")
|
||||
}
|
||||
let r =get_row()
|
||||
let c = get_col()
|
||||
let M = Matrix1(r, c)
|
||||
for (ir in 0..r) {
|
||||
for (ic in 0..c) {
|
||||
let v = this.get(ir, ic).getOrThrow()
|
||||
let vs = S.get(ir, ic).getOrThrow()
|
||||
M.set(ir, ic, v - vs)
|
||||
}
|
||||
}
|
||||
|
||||
return M
|
||||
}
|
||||
|
||||
public operator func *(a: FP32) { /* ---- 矩阵广播乘法 ---- */
|
||||
let r = get_row()
|
||||
let c = get_col()
|
||||
let M = Matrix1(r, c)
|
||||
for (ir in 0..r) {
|
||||
for (ic in 0..c){
|
||||
let v = this.get(ir, ic).getOrThrow()
|
||||
M.set(ir, ic, v * a)
|
||||
}
|
||||
}
|
||||
return M
|
||||
}
|
||||
|
||||
public operator func *(S: Matrix1) { /* ---- 逐元素相乘 ---- */
|
||||
if (!has_same_shape(S)) {
|
||||
throw IllegalArgumentException("Matrixes doesn't match")
|
||||
}
|
||||
let r =get_row()
|
||||
let c = get_col()
|
||||
let M = Matrix1(r, c)
|
||||
for (ir in 0..r) {
|
||||
for (ic in 0..c) {
|
||||
let v = this.get(ir, ic).getOrThrow()
|
||||
let vs = S.get(ir, ic).getOrThrow()
|
||||
M.set(ir, ic, v * vs)
|
||||
}
|
||||
}
|
||||
|
||||
return M
|
||||
}
|
||||
|
||||
public operator func /(a: FP32) { /* ---- 矩阵广播除法 ---- */
|
||||
let r = get_row()
|
||||
let c = get_col()
|
||||
let M = Matrix1(r, c)
|
||||
for (ir in 0..r) {
|
||||
for (ic in 0..c){
|
||||
let v = this.get(ir, ic).getOrThrow()
|
||||
M.set(ir, ic, v / a)
|
||||
}
|
||||
}
|
||||
return M
|
||||
}
|
||||
|
||||
public operator func /(S: Matrix1) { /* ---- 逐元素除法 ---- */
|
||||
if (!has_same_shape(S)) {
|
||||
throw IllegalArgumentException("Matrixes doesn't match")
|
||||
}
|
||||
let r =get_row()
|
||||
let c = get_col()
|
||||
let M = Matrix1(r, c)
|
||||
for (ir in 0..r) {
|
||||
for (ic in 0..c) {
|
||||
let v = this.get(ir, ic).getOrThrow()
|
||||
let vs = S.get(ir, ic).getOrThrow()
|
||||
M.set(ir, ic, v / vs)
|
||||
}
|
||||
}
|
||||
|
||||
return M
|
||||
}
|
||||
|
||||
/* ------------------------------ */
|
||||
/* -------- 矩阵点积乘法 -------- */
|
||||
public func dot(S: Matrix1) {
|
||||
if (!has_dot_shape(S)) {
|
||||
throw IllegalArgumentException("Matrixes doesn't match")
|
||||
}
|
||||
let r =get_row()
|
||||
let c = S.get_col()
|
||||
let M = Matrix1(r, c)
|
||||
for (ir in 0..get_row()) {
|
||||
for (ic in 0..S.get_col()) {
|
||||
var sum = FP32(0.0f32)
|
||||
for (ik in 0..get_col()) {
|
||||
let a = get(ir, ik).getOrThrow()
|
||||
let b = S.get(ik, ic).getOrThrow()
|
||||
sum += a * b
|
||||
}
|
||||
M.set(ir, ic, sum)
|
||||
}
|
||||
}
|
||||
return M
|
||||
}
|
||||
|
||||
/* -------- 克隆一个一模一样的矩阵 -------- */
|
||||
public func clone() {
|
||||
let r = get_row()
|
||||
let c = get_col()
|
||||
let N = Matrix1(r, c)
|
||||
for (ir in 0..r) {
|
||||
for (ic in 0..c) {
|
||||
let v = get(ir, ic).getOrThrow()
|
||||
N.set(ir, ic, v)
|
||||
}
|
||||
}
|
||||
return N
|
||||
}
|
||||
|
||||
/* -------- 矩阵转置,返回转置后的矩阵 -------- */
|
||||
public func T() {
|
||||
let r = get_row()
|
||||
let c = get_col()
|
||||
let N = Matrix1(c, r)
|
||||
for (ir in 0..r) {
|
||||
for (ic in 0..c) {
|
||||
let v = get(ir, ic).getOrThrow()
|
||||
N.set(ic, ir, v)
|
||||
}
|
||||
}
|
||||
return N
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
/* -------- 旧版的矩阵类,依赖c++编译的共享动态库 --------
|
||||
* TODO: 将矩阵元数据存储在仓颉部分(done at 202605xx ) */
|
||||
//public class Matrix {
|
||||
//
|
||||
// /* -------- 主初始化器,初始化并获取资源 -------- */
|
||||
// public init(row: UInt64, col: UInt64) {
|
||||
// unsafe { M = createMatrix(row, col) }
|
||||
// }
|
||||
//
|
||||
// /* -------- 初始化器,利用已经存在的元数据进行构建 -------- */
|
||||
// public Matrix(m: CPointer<Unit>) { M = m }
|
||||
//
|
||||
// /* -------- Xavier正态分布初始化 -------- */
|
||||
// public static func initWithXavierNormal(fan_in: UInt64, fan_out: UInt64) {
|
||||
// unsafe {
|
||||
// let m = xavierNormalInit(fan_in, fan_out);
|
||||
// return Matrix(m)
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// /* -------- Glorot均匀初始化 -------- */
|
||||
// public static func initWithXavierUniform(fan_in: UInt64, fan_out: UInt64) {
|
||||
// unsafe {
|
||||
// let m = xavierUniformInit(fan_in, fan_out);
|
||||
// return Matrix(m)
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// /* ======== 在析构函数中释放资源 ======== */
|
||||
// ~init() {
|
||||
// unsafe { freeMatrix(M) }
|
||||
// }
|
||||
//
|
||||
// /* -------- 取下标,由于不能返回值类型的引用被弃用 -------- */
|
||||
// //public operator func [](r: UInt64, c: UInt64): CPointer<Float32> {
|
||||
// // unsafe { return witchMatrix(M, r, c) }
|
||||
// //}
|
||||
//
|
||||
// /* -------- 获得某个矩阵单元的值 -------- */
|
||||
// public func get(r: UInt64, c: UInt64): Float32 {
|
||||
// unsafe { return getMatrix(M, r, c) }
|
||||
// }
|
||||
//
|
||||
// /* -------- 设置某个矩阵单元的值 --------*/
|
||||
// public func set(r: UInt64, c: UInt64, v: Float32) {
|
||||
// unsafe { setMatrix(M, r, c, v) }
|
||||
// }
|
||||
//
|
||||
// /* ------- 两个矩阵形状是否相同 --------
|
||||
// * FIXME: 这个函数所重载的运算符不贴切 */
|
||||
// public operator func ==(s: Matrix): Bool {
|
||||
// unsafe { return sameShapeMatrix(M, s.getData()) }
|
||||
// }
|
||||
//
|
||||
// /* -------- 获得行数 -------- */
|
||||
// public func getRow(): UInt64 {
|
||||
// unsafe { return getrowMatrix(M) }
|
||||
// }
|
||||
//
|
||||
// /* -------- 获取列数 --------*/
|
||||
// public func getCol(): UInt64 {
|
||||
// unsafe { return getcolMatrix(M) }
|
||||
// }
|
||||
//
|
||||
// /* -------- 获取指向元数据的C风格指针 -------- */
|
||||
// public func getData(): CPointer<Unit> { return M; }
|
||||
//
|
||||
// /* -------- 是否两个矩阵具有点积和的形状 -------- */
|
||||
// public func hasDotShapeWith(s: Matrix): Bool {
|
||||
// unsafe { return dotShapeMatrix(M, s.getData()) }
|
||||
// }
|
||||
//
|
||||
// /* -------- 逐元素算数运算 -------- */
|
||||
// /* -------------------------------- */
|
||||
// public operator func +(s: Matrix): Matrix {
|
||||
// unsafe {
|
||||
// if (!sameShapeMatrix(M, s.getData())) {
|
||||
// println('------------')
|
||||
// print_shape()
|
||||
// s.print_shape()
|
||||
// println('------------')
|
||||
// throw IllegalArgumentException("not same shape")
|
||||
// }
|
||||
// let r = plusMatrix(M, s.getData())
|
||||
// return Matrix(r)
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// public operator func +(t: Float32): Matrix {
|
||||
// unsafe {
|
||||
// let r = plusBroadcastMatrix(M, t)
|
||||
// return Matrix(r);
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// public operator func -(s: Matrix): Matrix {
|
||||
// unsafe {
|
||||
// if (!(this == s)) {
|
||||
// println('------------')
|
||||
// print_shape()
|
||||
// s.print_shape()
|
||||
// println('------------')
|
||||
// throw IllegalArgumentException("not same shape")
|
||||
// }
|
||||
// let r = subMatrix(M, s.getData())
|
||||
// return Matrix(r)
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// public operator func -(t: Float32): Matrix {
|
||||
// unsafe {
|
||||
// let r = subBroadcastMatrix(M, t)
|
||||
// return Matrix(r);
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// public operator func *(s: Matrix): Matrix {
|
||||
// unsafe {
|
||||
// if (!(this == s)) {
|
||||
// println('------------')
|
||||
// print_shape()
|
||||
// s.print_shape()
|
||||
// println('------------')
|
||||
// throw IllegalArgumentException("not same shape")
|
||||
// }
|
||||
// let r = multiMatrix(M, s.getData())
|
||||
// return Matrix(r)
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// public operator func *(t: Float32): Matrix {
|
||||
// unsafe {
|
||||
// let r = multiBroadcastMatrix(M, t)
|
||||
// return Matrix(r)
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// public operator func /(s: Matrix): Matrix {
|
||||
// unsafe {
|
||||
// if (!(this == s)) {
|
||||
// println('------------')
|
||||
// print_shape()
|
||||
// s.print_shape()
|
||||
// println('------------')
|
||||
// throw IllegalArgumentException("not same shape")
|
||||
// }
|
||||
// let r = divMatrix(M, s.getData())
|
||||
// return Matrix(r)
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// public operator func /(t: Float32): Matrix {
|
||||
// unsafe {
|
||||
// let r = divBroadcastMatrix(M, t)
|
||||
// return Matrix(r)
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// /* -------------------------------- */
|
||||
// /* -------------------------------- */
|
||||
//
|
||||
//
|
||||
// /* -------- 矩阵点积和算法 -------- */
|
||||
// public func dot (s: Matrix): Matrix {
|
||||
// unsafe {
|
||||
// if (!hasDotShapeWith(s)) {
|
||||
// println('------------')
|
||||
// print_shape()
|
||||
// s.print_shape()
|
||||
// println('------------')
|
||||
// throw IllegalArgumentException("not dot shape")
|
||||
// }
|
||||
// let r = dotMatrix(M, s.getData())
|
||||
// return Matrix(r)
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// /* -------- 矩阵转置 -------- */
|
||||
// public func T(): Matrix {
|
||||
// unsafe {
|
||||
// let r = transposeMatrix(M)
|
||||
// return Matrix(r)
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// /* -------- 打印矩阵 -------- */
|
||||
// public func print(): Unit {
|
||||
// unsafe {
|
||||
// for (i in 0..getRow()) {
|
||||
// for (j in 0..getCol()) {
|
||||
// let a = this.get(i, j)
|
||||
// print(a)
|
||||
// print(" ")
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
// println()
|
||||
// }
|
||||
//
|
||||
// /* -------- 打印形状 -------- */
|
||||
// public func print_shape() {
|
||||
// println("${getRow()}, ${getCol()}")
|
||||
// }
|
||||
//
|
||||
// /* -------- 保存矩阵 -------- */
|
||||
// public func save(file: File) {
|
||||
// unsafe {
|
||||
// for (i in 0..getRow()) {
|
||||
// for (j in 0..getCol()) {
|
||||
// file.write(this.get(i, j).toString().toArray())
|
||||
// if (j < getCol()-1) {
|
||||
// file.write(" ".toArray()) }
|
||||
// else { file.write("\n".toArray()) }
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// /* -------- 静态方法,加载单独项 -------- */
|
||||
// private static func loadItem(sr: StringReader<InputStream>) {
|
||||
// let s = sr.readUntil({r => r==r' ' || r ==r'\n'})
|
||||
// return Float32.parse(s.getOrThrow())
|
||||
// }
|
||||
//
|
||||
// /* -------- 静态方法,加载矩阵 -------- */
|
||||
// public static func load(sr: StringReader<InputStream>, row: UInt64, col: UInt64): Matrix {
|
||||
// let m = Matrix(row, col)
|
||||
// unsafe {
|
||||
// for (i in 0..m.getRow()) {
|
||||
// for (j in 0..m.getCol()) {
|
||||
// let f = loadItem(sr)
|
||||
// m.set(i, j, f)
|
||||
// }
|
||||
// }
|
||||
// m
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// /* -------- 矩阵的C风格指针 -------- */
|
||||
// private var M: CPointer<Unit>
|
||||
//}
|
||||
@@ -1,56 +0,0 @@
|
||||
package bot4cj.math
|
||||
import std.collection.*
|
||||
|
||||
/* -------- 张量 --------
|
||||
* INFO(20260524-qing): 输入一个维度列表 */
|
||||
class Tensor<T> {
|
||||
private let data: Array<T>
|
||||
private let dims: Array<Int64>
|
||||
private let size: Int64
|
||||
public init(dims: Array<Int64>, val: T) {
|
||||
var size = 1
|
||||
for (item in dims) { /* ---- 计算总量 ---- */
|
||||
size *= item
|
||||
}
|
||||
this.data = Array<T>(size, {_ => val})
|
||||
this.dims = dims
|
||||
this.size = size
|
||||
}
|
||||
|
||||
/* -------- 私有构造函数,用于内部创建张量 -------- */
|
||||
private init(dims: Array<Int64>, data: Array<T>) {
|
||||
this.dims = dims
|
||||
this.data = data
|
||||
var size = 1
|
||||
for (item in dims) {
|
||||
size *= item
|
||||
}
|
||||
this.size = size
|
||||
}
|
||||
|
||||
/* -------- 获取维度列表(形状) -------- */
|
||||
public func shape(): Array<Int64> {
|
||||
return dims
|
||||
}
|
||||
|
||||
/* -------- 逐元素进行某种运算 -------- */
|
||||
public func oneByOne(S: Tensor<T>, f: (T, T) -> T): Tensor<T> {
|
||||
let s1 = S.shape()
|
||||
|
||||
/* ---- 比较维度数量 ---- */
|
||||
if (s1 != dims) {
|
||||
throw IllegalArgumentException("Tensors' shape are not same")
|
||||
}
|
||||
|
||||
/* ---- 比较维度形状 ----*/
|
||||
for (i in 0..dims.size) {
|
||||
if (s1[i] != dims[i]) {
|
||||
throw IllegalArgumentException("Tensors' shape are not same")
|
||||
}
|
||||
}
|
||||
|
||||
/* ---- 逐个进行计算 ---- */
|
||||
let res_data = Array<T>(size, {i => f(data[i], S.data[i])})
|
||||
return Tensor<T>(dims, res_data)
|
||||
}
|
||||
}
|
||||
@@ -1,379 +0,0 @@
|
||||
// cnn.cj - 修正 forward 中的偏置加法
|
||||
package bot4cj.math
|
||||
|
||||
import std.math.*
|
||||
import std.random.*
|
||||
import std.collection.*
|
||||
|
||||
public class ConvLayer {
|
||||
|
||||
let nInputChannels: Int64
|
||||
let nOutputChannels: Int64
|
||||
let kernelHeight: Int64
|
||||
let kernelWidth: Int64
|
||||
let stride: Int64
|
||||
let padding: Int64
|
||||
let learningRate: FP32
|
||||
let f: Af
|
||||
|
||||
var connectingMatrix: Matrix1
|
||||
var kernels: Array<Matrix1>
|
||||
var biases: Matrix1
|
||||
|
||||
var prevInput: Option<Matrix1> = Option.None
|
||||
var prevInputShape: Option<(Int64, Int64, Int64, Int64)> = Option.None
|
||||
var prevCombination: Option<Matrix1> = Option.None
|
||||
var prevOutput: Option<Matrix1> = Option.None
|
||||
var prevOutputShape: Option<(Int64, Int64, Int64, Int64)> = Option.None
|
||||
|
||||
var kernelsGrad: Option<Array<Matrix1>> = Option.None
|
||||
var biasGrad: Option<Matrix1> = Option.None
|
||||
|
||||
public init(
|
||||
inCh: Int64,
|
||||
outCh: Int64,
|
||||
kernelSize: (Int64, Int64),
|
||||
stride: Int64,
|
||||
padding: Int64,
|
||||
lr: FP32,
|
||||
activation: Af,
|
||||
connMatrix: ?Matrix1
|
||||
) {
|
||||
nInputChannels = inCh
|
||||
nOutputChannels = outCh
|
||||
kernelHeight = kernelSize[0]
|
||||
kernelWidth = kernelSize[1]
|
||||
this.stride = stride
|
||||
this.padding = padding
|
||||
learningRate = lr
|
||||
this.f = activation
|
||||
|
||||
// 初始化 connectingMatrix
|
||||
connectingMatrix = match (connMatrix) {
|
||||
case m: Matrix1 => m
|
||||
case _ => Matrix1(nInputChannels, nOutputChannels)
|
||||
}
|
||||
let isNone = match (connMatrix) {
|
||||
case _: Matrix1 => false
|
||||
case _ => true
|
||||
}
|
||||
if (isNone) {
|
||||
for (r in 0..nInputChannels) {
|
||||
for (c in 0..nOutputChannels) {
|
||||
connectingMatrix.set(r, c, FP32(1.0))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 初始化 kernels
|
||||
kernels = Array(nOutputChannels, { _ => Matrix1(1, 1) })
|
||||
for (oc in 0..nOutputChannels) {
|
||||
var nConnections = 0
|
||||
for (ic in 0..nInputChannels) {
|
||||
let val = connectingMatrix.get(ic, oc).getOrThrow()
|
||||
if (val.getValue() > 0.0) {
|
||||
nConnections += 1
|
||||
}
|
||||
}
|
||||
let kernelLen = nConnections * kernelHeight * kernelWidth
|
||||
let kernel = Matrix1(1, kernelLen)
|
||||
|
||||
Ini.xavierUniformInit(kernel, 1, kernelLen)
|
||||
|
||||
kernels[oc] = kernel
|
||||
}
|
||||
|
||||
// 初始化 biases
|
||||
biases = Matrix1(1, nOutputChannels)
|
||||
Ini.xavierUniformInit(biases, 1, nOutputChannels)
|
||||
}
|
||||
|
||||
// ==================== 前向传播 ====================
|
||||
public func forward(x: Matrix1, inputShape: (Int64, Int64, Int64, Int64)): Matrix1 {
|
||||
prevInput = Option.Some(x)
|
||||
prevInputShape = Option.Some(inputShape)
|
||||
|
||||
let batch = inputShape[0]
|
||||
let inH = inputShape[1]
|
||||
let inW = inputShape[2]
|
||||
|
||||
let outH = (inH + 2 * padding - kernelHeight) / stride + 1
|
||||
let outW = (inW + 2 * padding - kernelWidth) / stride + 1
|
||||
let outSize = batch * outH * outW
|
||||
|
||||
var combination = Matrix1(outSize, nOutputChannels)
|
||||
var output = Matrix1(outSize, nOutputChannels)
|
||||
|
||||
for (b in 0..batch) {
|
||||
for (oh in 0..outH) {
|
||||
for (ow in 0..outW) {
|
||||
let outRow = b * outH * outW + oh * outW + ow
|
||||
|
||||
for (oc in 0..nOutputChannels) {
|
||||
var sum = FP32(0.0)
|
||||
var kernelIdx = 0
|
||||
|
||||
for (ic in 0..nInputChannels) {
|
||||
let connVal = connectingMatrix.get(ic, oc).getOrThrow()
|
||||
if (connVal.getValue() > 0.0) {
|
||||
for (kh in 0..kernelHeight) {
|
||||
for (kw in 0..kernelWidth) {
|
||||
let ih = oh * stride + kh - padding
|
||||
let iw = ow * stride + kw - padding
|
||||
|
||||
if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
|
||||
let inRow = b * inH * inW + ih * inW + iw
|
||||
let xVal = x.get(inRow, ic).getOrThrow()
|
||||
let kVal = kernels[oc].get(0, kernelIdx).getOrThrow()
|
||||
sum = sum + xVal * kVal
|
||||
}
|
||||
kernelIdx += 1
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
combination.set(outRow, oc, sum)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 手动广播偏置加法(因为 Matrix1 的 + 要求形状相同)
|
||||
for (r in 0..outSize) {
|
||||
for (c in 0..nOutputChannels) {
|
||||
let val = combination.get(r, c).getOrThrow()
|
||||
let biasVal = biases.get(0, c).getOrThrow()
|
||||
combination.set(r, c, val + biasVal)
|
||||
}
|
||||
}
|
||||
|
||||
prevCombination = Option.Some(combination)
|
||||
|
||||
output = f.forward(combination)
|
||||
prevOutput = Option.Some(output)
|
||||
prevOutputShape = Option.Some((batch, outH, outW, nOutputChannels))
|
||||
|
||||
return output
|
||||
}
|
||||
|
||||
// ==================== 反向传播 ====================
|
||||
public func backward(errs: Matrix1): Matrix1 {
|
||||
var x = Matrix1(1, 1)
|
||||
match (prevInput) {
|
||||
case Some(T) => x = T
|
||||
case _ => throw Exception("No input data")
|
||||
}
|
||||
|
||||
var shape = (0, 0, 0, 0)
|
||||
match (prevInputShape) {
|
||||
case Some(T) => shape = T
|
||||
case _ => throw Exception("No input shape")
|
||||
}
|
||||
|
||||
var combination = Matrix1(1, 1)
|
||||
match (prevCombination) {
|
||||
case Some(val) => combination = val
|
||||
case _ => throw Exception("No combination")
|
||||
}
|
||||
|
||||
var outShape = (0, 0, 0, 0)
|
||||
match (prevOutputShape) {
|
||||
case Some(val) => outShape = val
|
||||
case _ => throw Exception("No output shape")
|
||||
}
|
||||
|
||||
let batch = shape[0]
|
||||
let inH = shape[1]
|
||||
let inW = shape[2]
|
||||
let outH = outShape[1]
|
||||
let outW = outShape[2]
|
||||
|
||||
let delta = errs * f.backward(combination)
|
||||
|
||||
calcKernelGradient(x, delta, batch, inH, inW, outH, outW)
|
||||
calcBiasGradient(delta)
|
||||
|
||||
let errsUp = calcUpsampledError(delta, batch, inH, inW, outH, outW)
|
||||
|
||||
return errsUp
|
||||
}
|
||||
|
||||
/**
|
||||
* 计算卷积核梯度
|
||||
*/
|
||||
private func calcKernelGradient(
|
||||
x: Matrix1,
|
||||
delta: Matrix1,
|
||||
batch: Int64,
|
||||
inH: Int64,
|
||||
inW: Int64,
|
||||
outH: Int64,
|
||||
outW: Int64
|
||||
): Unit {
|
||||
var grads = Array(nOutputChannels, { _ => Matrix1(1, 1) })
|
||||
|
||||
for (oc in 0..nOutputChannels) {
|
||||
let kernelCols = kernels[oc].get_col()
|
||||
var grad = Matrix1(1, kernelCols)
|
||||
|
||||
for (b in 0..batch) {
|
||||
for (oh in 0..outH) {
|
||||
for (ow in 0..outW) {
|
||||
let deltaRow = b * outH * outW + oh * outW + ow
|
||||
let deltaVal = delta.get(deltaRow, oc).getOrThrow()
|
||||
|
||||
var kernelIdx = 0
|
||||
for (ic in 0..nInputChannels) {
|
||||
let connVal = connectingMatrix.get(ic, oc).getOrThrow()
|
||||
if (connVal.getValue() > 0.0) {
|
||||
for (kh in 0..kernelHeight) {
|
||||
for (kw in 0..kernelWidth) {
|
||||
let ih = oh * stride + kh - padding
|
||||
let iw = ow * stride + kw - padding
|
||||
|
||||
if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
|
||||
let inRow = b * inH * inW + ih * inW + iw
|
||||
let xVal = x.get(inRow, ic).getOrThrow()
|
||||
let oldGrad = grad.get(0, kernelIdx).getOrThrow()
|
||||
grad.set(0, kernelIdx, oldGrad + xVal * deltaVal)
|
||||
}
|
||||
kernelIdx += 1
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
grads[oc] = grad
|
||||
}
|
||||
|
||||
kernelsGrad = Option.Some(grads)
|
||||
}
|
||||
|
||||
/**
|
||||
* 计算偏置梯度
|
||||
*/
|
||||
private func calcBiasGradient(delta: Matrix1): Unit {
|
||||
let rows = delta.get_row()
|
||||
var grad = Matrix1(1, nOutputChannels)
|
||||
|
||||
for (oc in 0..nOutputChannels) {
|
||||
var sum = FP32(0.0)
|
||||
for (r in 0..rows) {
|
||||
sum = sum + delta.get(r, oc).getOrThrow()
|
||||
}
|
||||
grad.set(0, oc, sum)
|
||||
}
|
||||
|
||||
biasGrad = Option.Some(grad)
|
||||
}
|
||||
|
||||
/**
|
||||
* 计算传递到上一层的误差
|
||||
*/
|
||||
private func calcUpsampledError(
|
||||
delta: Matrix1,
|
||||
batch: Int64,
|
||||
inH: Int64,
|
||||
inW: Int64,
|
||||
outH: Int64,
|
||||
outW: Int64
|
||||
): Matrix1 {
|
||||
let inSize = batch * inH * inW
|
||||
var errsUp = Matrix1(inSize, nInputChannels)
|
||||
|
||||
for (b in 0..batch) {
|
||||
for (oh in 0..outH) {
|
||||
for (ow in 0..outW) {
|
||||
let deltaRow = b * outH * outW + oh * outW + ow
|
||||
|
||||
for (oc in 0..nOutputChannels) {
|
||||
let deltaVal = delta.get(deltaRow, oc).getOrThrow()
|
||||
|
||||
var kernelIdx = 0
|
||||
for (ic in 0..nInputChannels) {
|
||||
let connVal = connectingMatrix.get(ic, oc).getOrThrow()
|
||||
if (connVal.getValue() > 0.0) {
|
||||
for (kh in 0..kernelHeight) {
|
||||
for (kw in 0..kernelWidth) {
|
||||
let ih = oh * stride + kh - padding
|
||||
let iw = ow * stride + kw - padding
|
||||
|
||||
if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
|
||||
let inRow = b * inH * inW + ih * inW + iw
|
||||
let kVal = kernels[oc].get(0, kernelIdx).getOrThrow()
|
||||
let oldErr = errsUp.get(inRow, ic).getOrThrow()
|
||||
errsUp.set(inRow, ic, oldErr + deltaVal * kVal)
|
||||
}
|
||||
kernelIdx += 1
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return errsUp
|
||||
}
|
||||
|
||||
// ==================== 参数更新 ====================
|
||||
public func update(discount: Float32): Unit {
|
||||
var kGrad = Array(0, { _ => Matrix1(1, 1) })
|
||||
match (kernelsGrad) {
|
||||
case Some(val) => kGrad = val
|
||||
case _ => throw Exception("No kernel gradients")
|
||||
}
|
||||
|
||||
var bGrad = Matrix1(1, 1)
|
||||
match (biasGrad) {
|
||||
case Some(val) => bGrad = val
|
||||
case _ => throw Exception("No bias gradient")
|
||||
}
|
||||
|
||||
for (oc in 0..nOutputChannels) {
|
||||
let rows = kernels[oc].get_row()
|
||||
let cols = kernels[oc].get_col()
|
||||
for (r in 0..rows) {
|
||||
for (c in 0..cols) {
|
||||
let oldVal = kernels[oc].get(r, c).getOrThrow()
|
||||
let gradVal = kGrad[oc].get(r, c).getOrThrow()
|
||||
kernels[oc].set(r, c, oldVal - gradVal * learningRate * FP32(discount))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (oc in 0..nOutputChannels) {
|
||||
let oldBias = biases.get(0, oc).getOrThrow()
|
||||
let gradBias = bGrad.get(0, oc).getOrThrow()
|
||||
biases.set(0, oc, oldBias - gradBias * learningRate * FP32(discount))
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== 形状访问器 ====================
|
||||
public func getInShape(): (Int64, Int64) {
|
||||
(nInputChannels, kernelHeight * kernelWidth)
|
||||
}
|
||||
|
||||
public func getOutShape(): Int64 {
|
||||
nOutputChannels
|
||||
}
|
||||
|
||||
public func getLR(): FP32 {
|
||||
learningRate
|
||||
}
|
||||
|
||||
public func getFType(): String {
|
||||
f.get_type()
|
||||
}
|
||||
|
||||
public func calculateOutputSize(inputHeight: Int64, inputWidth: Int64): (Int64, Int64) {
|
||||
let outH = (inputHeight + 2 * padding - kernelHeight) / stride + 1
|
||||
let outW = (inputWidth + 2 * padding - kernelWidth) / stride + 1
|
||||
(outH, outW)
|
||||
}
|
||||
}
|
||||
@@ -1,11 +0,0 @@
|
||||
package bot4cj.math
|
||||
//foreign func start_fft(M: CPointer<Unit>): CPointer<Unit>
|
||||
//foreign func gen_wave(): CPointer<Unit>
|
||||
/*-------- 生成测试用的正弦波 --------*/
|
||||
//public func genWave(): Matrix {
|
||||
// unsafe { return Matrix(gen_wave()) }
|
||||
//}
|
||||
/*-------- 执行fft算法 --------*/
|
||||
//public func start_fft(M: Matrix): Matrix {
|
||||
// unsafe { return Matrix(start_fft(M.getData())) }
|
||||
//}
|
||||
@@ -1,136 +0,0 @@
|
||||
// flatten.cj - 移除调试代码
|
||||
package bot4cj.math
|
||||
|
||||
import std.math.*
|
||||
import std.collection.*
|
||||
|
||||
public class FlattenLayer {
|
||||
|
||||
let nInputChannels: Int64
|
||||
let inputHeight: Int64
|
||||
let inputWidth: Int64
|
||||
let outputSize: Int64
|
||||
|
||||
var prevInput: Option<Matrix1>
|
||||
var prevInputShape: Option<(Int64, Int64, Int64, Int64)>
|
||||
var prevOutput: Option<Matrix1>
|
||||
var prevOutputShape: Option<(Int64, Int64)>
|
||||
|
||||
public init(inputChannels: Int64, inputHeight: Int64, inputWidth: Int64) {
|
||||
nInputChannels = inputChannels
|
||||
this.inputHeight = inputHeight
|
||||
this.inputWidth = inputWidth
|
||||
outputSize = inputChannels * inputHeight * inputWidth
|
||||
|
||||
prevInput = None
|
||||
prevInputShape = None
|
||||
prevOutput = None
|
||||
prevOutputShape = None
|
||||
}
|
||||
|
||||
public init(inputShape: (Int64, Int64, Int64)) {
|
||||
nInputChannels = inputShape[0]
|
||||
inputHeight = inputShape[1]
|
||||
inputWidth = inputShape[2]
|
||||
outputSize = nInputChannels * inputHeight * inputWidth
|
||||
|
||||
prevInput = None
|
||||
prevInputShape = None
|
||||
prevOutput = None
|
||||
prevOutputShape = None
|
||||
}
|
||||
|
||||
public func forward(x: Matrix1, inputShape: (Int64, Int64, Int64, Int64)): Matrix1 {
|
||||
prevInput = x
|
||||
prevInputShape = inputShape
|
||||
|
||||
let batch = inputShape[0]
|
||||
let inH = inputShape[1]
|
||||
let inW = inputShape[2]
|
||||
let inCh = inputShape[3]
|
||||
|
||||
let totalRows = batch * outputSize
|
||||
var output = Matrix1(totalRows, 1)
|
||||
|
||||
for (b in 0..batch) {
|
||||
for (h in 0..inH) {
|
||||
for (w in 0..inW) {
|
||||
for (ch in 0..inCh) {
|
||||
let inRow = b * inH * inW + h * inW + w
|
||||
let outRow = b * outputSize + (ch * inH * inW + h * inW + w)
|
||||
if (outRow < totalRows) {
|
||||
let val = x.get(inRow, ch).getOrThrow()
|
||||
output.set(outRow, 0, val)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
prevOutput = output
|
||||
prevOutputShape = (totalRows, 1)
|
||||
|
||||
return output
|
||||
}
|
||||
|
||||
public func backward(errs: Matrix1): Matrix1 {
|
||||
var shape = (0, 0, 0, 0)
|
||||
match (prevInputShape) {
|
||||
case Some(T) => shape = T
|
||||
case _ => throw Exception("FlattenLayer.backward: No input shape. Call forward() first.")
|
||||
}
|
||||
|
||||
let batch = shape[0]
|
||||
let inH = shape[1]
|
||||
let inW = shape[2]
|
||||
let inCh = shape[3]
|
||||
|
||||
let inSize = batch * inH * inW
|
||||
var errsUp = Matrix1(inSize, nInputChannels)
|
||||
|
||||
let errsCols = errs.get_col()
|
||||
|
||||
for (b in 0..batch) {
|
||||
for (h in 0..inH) {
|
||||
for (w in 0..inW) {
|
||||
let inRow = b * inH * inW + h * inW + w
|
||||
|
||||
for (ch in 0..inCh) {
|
||||
let outRow = b * outputSize + (ch * inH * inW + h * inW + w)
|
||||
|
||||
if (outRow < errs.get_row()) {
|
||||
if (errsCols == 1) {
|
||||
let val = errs.get(outRow, 0).getOrThrow()
|
||||
errsUp.set(inRow, ch, val)
|
||||
} else {
|
||||
var sum = FP32(0.0)
|
||||
for (col in 0..errsCols) {
|
||||
sum = sum + errs.get(outRow, col).getOrThrow()
|
||||
}
|
||||
errsUp.set(inRow, ch, sum)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return errsUp
|
||||
}
|
||||
|
||||
public func update(discount: Float32): Unit {
|
||||
// 无参数
|
||||
}
|
||||
|
||||
public func getInShape(): (Int64, Int64, Int64) {
|
||||
(nInputChannels, inputHeight, inputWidth)
|
||||
}
|
||||
|
||||
public func getOutShape(): Int64 {
|
||||
outputSize
|
||||
}
|
||||
|
||||
public func getOutputSize(): Int64 {
|
||||
outputSize
|
||||
}
|
||||
}
|
||||
@@ -1,331 +0,0 @@
|
||||
package bot4cj.math
|
||||
import std.fs.*
|
||||
import std.io.*
|
||||
import std.convert.*
|
||||
//foreign func Sigmoid_do(m: CPointer<Unit>): CPointer<Unit>
|
||||
//foreign func Sigmoid_dot(m: CPointer<Unit>): CPointer<Unit>
|
||||
//foreign func ReLU_do(m: CPointer<Unit>): CPointer<Unit>
|
||||
//foreign func ReLU_dot(m: CPointer<Unit>): CPointer<Unit>
|
||||
//foreign func Leaky_ReLU_do(m: CPointer<Unit>, alpha: Float32): CPointer<Unit>
|
||||
//foreign func Leaky_ReLU_dot(m: CPointer<Unit>, alpha: Float32): CPointer<Unit>
|
||||
//foreign func Tanh_do(m: CPointer<Unit>): CPointer<Unit>
|
||||
//foreign func Tanh_dot(m: CPointer<Unit>): CPointer<Unit>
|
||||
|
||||
|
||||
/* -------- 新版的神经网络,使用原生矩阵对象实现 --------
|
||||
* TODO(20260517-qing): 增加batch批次处理,使用ArrayList<Matrix1>
|
||||
* TODO*(20260519-qing): 将神经网络转换为MatrixObject矩阵式对象
|
||||
* TODO(20260524-qing): 增加Adam优化器 */
|
||||
|
||||
public class Nn1 {
|
||||
|
||||
let inputno: Int64 /* -------- 输入 -------- */
|
||||
let outputno: Int64 /* -------- 输出 -------- */
|
||||
var weights: Matrix1 /* -- (inputs, outputs) - */
|
||||
var bias: Matrix1 /* ----- (1, output) ---- */
|
||||
let learning_rate: FP32 /* ------- 学习率 ------- */
|
||||
let f: Af /* ------ 激活函数 ------ */
|
||||
|
||||
/*------------------*/
|
||||
var prev_inputs: Option<Matrix1> = None /* ---- 之前的输入 ---- */
|
||||
var prev_raw_inputs: Option<Matrix1> = None /* ---- 之前的原始输出 ---- */
|
||||
var prev_outputs: Option<Matrix1> = None /* ---- 之前的输出 ---- */
|
||||
/*------------------*/
|
||||
var weights_grad: Option<Matrix1> = None /* ---- 权重梯度 ---- */
|
||||
var bias_grad: Option<Matrix1> = None /* ---- 偏置梯度 ---- */
|
||||
|
||||
|
||||
/* -------- 初始化器 --------
|
||||
* TODO(20260516-qing): 不要把初始化方法写死,使用枚举类 */
|
||||
public init(i: Int64, o: Int64, lr: Float32, f: Af) {
|
||||
inputno = i
|
||||
outputno =o
|
||||
learning_rate=FP32(lr)
|
||||
this.f = f
|
||||
weights = Matrix1(inputno, outputno)
|
||||
bias = Matrix1(1, outputno)
|
||||
|
||||
/* -------- 初始化权重 -------- */
|
||||
Ini.xavierUniformInit(weights, inputno, outputno)
|
||||
}
|
||||
|
||||
/* -------- 神经网络推理 -------- */
|
||||
public func forward(x: Matrix1) {
|
||||
prev_inputs = x
|
||||
var r = x.dot(weights)
|
||||
r = r + bias
|
||||
prev_raw_inputs = r
|
||||
let r1 = f.forward(r)
|
||||
prev_outputs = r1
|
||||
return r1
|
||||
}
|
||||
|
||||
/* -------- 反向传播算法 -------- */
|
||||
public func backward(errs: Matrix1)
|
||||
{
|
||||
if (is_softmax(f)) {
|
||||
return softmax_backward(errs)
|
||||
}
|
||||
let fbak = f.backward(prev_raw_inputs.getOrThrow())
|
||||
let first = errs * fbak /* 逐元素相乘 */
|
||||
let mx = prev_inputs.getOrThrow().T()
|
||||
weights_grad = mx.dot(first)
|
||||
bias_grad = first
|
||||
let errs_up = first.dot(weights.T())
|
||||
return errs_up
|
||||
}
|
||||
|
||||
/* -------- 更新权重值、偏爱因子 --------
|
||||
* NOTE: discount(折扣)是为了方便深度强化学习给出奖励值 */
|
||||
public func update(discount: Float32) {
|
||||
weights = weights - weights_grad.getOrThrow() * learning_rate * FP32(discount)
|
||||
bias = bias - bias_grad.getOrThrow() * learning_rate * FP32(discount)
|
||||
}
|
||||
|
||||
/**/
|
||||
private static func is_softmax(f: Af) {
|
||||
return match (f) {
|
||||
case Af.Softmax => true
|
||||
case _ => false
|
||||
}
|
||||
}
|
||||
|
||||
private func softmax_backward(errs: Matrix1) {
|
||||
/* 使用雅可比矩阵的方式 */
|
||||
let y = prev_outputs.getOrThrow() /* softmax 输出 */
|
||||
let dy = errs /* 上游梯度 ∂L/∂y */
|
||||
let jacobian = f.backward(y) /* n×n 雅可比矩阵 */
|
||||
|
||||
/* 将 dy 展平成列向量 n×1 */
|
||||
let row = dy.get_row()
|
||||
let col = dy.get_col()
|
||||
let n = row * col
|
||||
var dy_vec = Matrix1(n, 1)
|
||||
var idx = 0
|
||||
for (r in 0..row) {
|
||||
for (c in 0..col) {
|
||||
let val = dy.get(r, c).getOrThrow()
|
||||
dy_vec.set(idx, 0, val)
|
||||
idx = idx + 1
|
||||
}
|
||||
}
|
||||
|
||||
/* dx_vec = J · dy_vec (n×n × n×1 = n×1) */
|
||||
var dx_vec = Matrix1(n, 1)
|
||||
for (i in 0..n) {
|
||||
var sum = 0.0f32
|
||||
for (j in 0..n) {
|
||||
let j_val = jacobian.get(i, j).getOrThrow().getValue()
|
||||
let dy_val = dy_vec.get(j, 0).getOrThrow().getValue()
|
||||
sum = sum + j_val * dy_val
|
||||
}
|
||||
dx_vec.set(i, 0, FP32(sum))
|
||||
}
|
||||
|
||||
/* 恢复为原始形状,作为 first */
|
||||
var first = Matrix1(row, col)
|
||||
idx = 0
|
||||
for (r in 0..row) {
|
||||
for (c in 0..col) {
|
||||
let val = dx_vec.get(idx, 0).getOrThrow()
|
||||
first.set(r, c, val)
|
||||
idx = idx + 1
|
||||
}
|
||||
}
|
||||
|
||||
let mx = prev_inputs.getOrThrow().T()
|
||||
weights_grad = mx.dot(first)
|
||||
bias_grad = first
|
||||
let errs_up = first.dot(weights.T())
|
||||
return errs_up
|
||||
}
|
||||
|
||||
|
||||
/* -------- 形状访问器 -------- */
|
||||
public func get_in_shape() {
|
||||
return inputno }
|
||||
|
||||
public func get_out_shape() {
|
||||
return outputno }
|
||||
|
||||
public func get_lr() {
|
||||
return learning_rate }
|
||||
|
||||
public func get_f_type() {
|
||||
return f.get_type() }
|
||||
|
||||
public func get_f_val() {
|
||||
return f.get_value() }
|
||||
/* ---------------------------- */
|
||||
|
||||
}
|
||||
|
||||
|
||||
///* -------- 旧版激活函数 --------*/
|
||||
//public enum ActivationFuncs {
|
||||
// Sigmoid | ReLU | Leaky_ReLU | Tanh
|
||||
// func fdo(m: Matrix): Matrix { // 暂时只设置ReLU
|
||||
// unsafe {
|
||||
// match(this) {
|
||||
// case Sigmoid => Matrix(Sigmoid_do(m.getData()))
|
||||
// case ReLU => Matrix(ReLU_do(m.getData()))
|
||||
// case Leaky_ReLU => Matrix(Leaky_ReLU_do(m.getData(), 0.1f32))
|
||||
// case Tanh => Matrix(Tanh_do(m.getData()))
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
// func fdot(m: Matrix): Matrix {
|
||||
// unsafe {
|
||||
// match(this) {
|
||||
// case Sigmoid => Matrix(Sigmoid_dot(m.getData()))
|
||||
// case ReLU => Matrix(ReLU_dot(m.getData()))
|
||||
// case Leaky_ReLU => Matrix(Leaky_ReLU_dot(m.getData(), 0.1f32))
|
||||
// case Tanh => Matrix(Tanh_dot(m.getData()))
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
// func toString(): String {
|
||||
// match(this) {
|
||||
// case Sigmoid => "Sigmoid"
|
||||
// case ReLU => "ReLU"
|
||||
// case Leaky_ReLU => "Leaky_ReLU"
|
||||
// case Tanh => "Tanh"
|
||||
// }
|
||||
// }
|
||||
// static func parse(s: String) {
|
||||
// match(s) {
|
||||
// case "Sigmoid" => ActivationFuncs.Sigmoid
|
||||
// case "ReLU" => ActivationFuncs.ReLU
|
||||
// case "Leaky_ReLU" => ActivationFuncs.Leaky_ReLU
|
||||
// case "Tanh" => ActivationFuncs.Tanh
|
||||
// case _ => throw IllegalArgumentException()
|
||||
// }
|
||||
// }
|
||||
//}
|
||||
|
||||
///* -------- 旧版的神经网络类 --------
|
||||
// * 这个实现依赖C++语言的动态分享库,特点是速度快
|
||||
// * 这个实现可能是有问题的 */
|
||||
//public class Nn {
|
||||
//
|
||||
// /* -------- 主初始化器 -------- */
|
||||
// public init(
|
||||
// inputno: UInt64, outputno: UInt64, f: ActivationFuncs, learning_rate: Float32) {
|
||||
// this.inputno = inputno
|
||||
// this.outputno = outputno
|
||||
// weights = Matrix.initWithXavierUniform(inputno, outputno)
|
||||
// bias = Matrix(1, outputno)
|
||||
// this.f = f
|
||||
// this.learning_rate = learning_rate
|
||||
// }
|
||||
//
|
||||
// /* -------- 提供了权重矩阵的初始化器 -------- */
|
||||
// public init(
|
||||
// inputno: UInt64, outputno: UInt64, f: ActivationFuncs, learning_rate: Float32, weights: Matrix, bias: Matrix) {
|
||||
// this.inputno = inputno
|
||||
// this.outputno = outputno
|
||||
// this.weights = weights
|
||||
// this.bias = bias
|
||||
// this.f = f
|
||||
// this.learning_rate = learning_rate
|
||||
// }
|
||||
//
|
||||
// /* -------- 神经网络推理 -------- */
|
||||
// func forward(x: Matrix): Matrix {
|
||||
// /* x => (1, inputs) */
|
||||
// let r = x.dot(weights) + bias /* (1, outputno) */
|
||||
// return r
|
||||
// }
|
||||
//
|
||||
// /* -------- 实施激活函数 -------- */
|
||||
// func fdo(x: Matrix): Matrix {
|
||||
// return f.fdo(x) }
|
||||
//
|
||||
// /* -------- 激活函数导数 -------- */
|
||||
// func fdot(x: Matrix): Matrix {
|
||||
// return f.fdot(x) }
|
||||
//
|
||||
// /* This is Wrong */
|
||||
// //func bak_bsgrad(prev_raw_outputs_dotf: Matrix, errs: Matrix): Matrix {
|
||||
// /* prev_raw_outputs_dotf => (1, outputno)
|
||||
// * errs => (1, outputno) */
|
||||
//
|
||||
// // return errs * prev_raw_outputs_dotf
|
||||
// //}
|
||||
//
|
||||
// /* -------- 权重梯度的传播 -------- */
|
||||
// func bak_wsgrad(prev_inputs: Matrix, delta: Matrix) {
|
||||
// /* prev_inputs => (1, inputs) */
|
||||
// return prev_inputs.T().dot(delta) /* (inputs, outputs) */
|
||||
// }
|
||||
//
|
||||
// /* -------- 计算上一层的误差 -------- */
|
||||
// func bak_uperrs(delta: Matrix, prev_dotf: Matrix) {
|
||||
// let first = weights.dot(delta.T()).T() // (1, inputs)
|
||||
//
|
||||
// var uperrs = first * prev_dotf
|
||||
// uperrs
|
||||
// }
|
||||
//
|
||||
// /* -------- 更新权重、偏爱因子 -------- */
|
||||
// func update(weights_grad: Matrix, bias_grad: Matrix, discount: Float32) {
|
||||
// weights -= weights_grad * learning_rate * discount
|
||||
// bias -= bias_grad * learning_rate * discount
|
||||
// }
|
||||
//
|
||||
// /* -------- 打印权重、偏爱因子 -------- */
|
||||
// func print() {
|
||||
// weights.print()
|
||||
// bias.print()
|
||||
// }
|
||||
//
|
||||
// /* -------- 打印神经网络全连接层的形状 -------- */
|
||||
// func print_shape() {
|
||||
// println("${inputno}, ${outputno}") }
|
||||
//
|
||||
// /* -------- 储存神经网络 -------- */
|
||||
// func save(file: File) {
|
||||
// file.write("FullConnectLayer ".toArray())
|
||||
// file.write(inputno.toString().toArray())
|
||||
// file.write(" ".toArray())
|
||||
// file.write(outputno.toString().toArray())
|
||||
// file.write(" ".toArray())
|
||||
// file.write(learning_rate.toString().toArray())
|
||||
// file.write(" ".toArray())
|
||||
// file.write(f.toString().toArray())
|
||||
// file.write("\n".toArray())
|
||||
// weights.save(file)
|
||||
// //file.write("\n".toArray())
|
||||
// bias.save(file)
|
||||
// //file.write("\n".toArray())
|
||||
// }
|
||||
//
|
||||
// /* -------- 静态方法,加载单个项目 -------- */
|
||||
// static func loadItem<T>(sr: StringReader<InputStream>) where T <: Parsable<T>
|
||||
// {
|
||||
// let s = sr.readUntil({r => r == r" " || r == r"\n"})
|
||||
// let ss = s.getOrThrow().trimAscii()
|
||||
// println(ss)
|
||||
// return T.parse(ss)
|
||||
// }
|
||||
//
|
||||
// /* -------- 静态方法,加载神经网络 -------- */
|
||||
// static func load(sr: StringReader<InputStream>)
|
||||
// {
|
||||
// let inputno = loadItem<UInt64>(sr)
|
||||
// let outputno = loadItem<UInt64>(sr)
|
||||
// let lr = loadItem<Float32>(sr)
|
||||
// let fna = sr.readUntil(r'\n')
|
||||
// let f = ActivationFuncs.parse(fna.getOrThrow().trimAscii())
|
||||
// let weights = Matrix.load(sr, inputno, outputno)
|
||||
// let bias = Matrix.load(sr, 1, outputno)
|
||||
// Nn(inputno, outputno, f, lr, weights, bias)
|
||||
// }
|
||||
//
|
||||
// let inputno: UInt64
|
||||
// let outputno: UInt64
|
||||
// var weights: Matrix /* (inputs, outputs) */
|
||||
// var bias: Matrix /* (1, outputs) */
|
||||
// var learning_rate: Float32
|
||||
// let f: ActivationFuncs
|
||||
//}
|
||||
@@ -1,164 +0,0 @@
|
||||
package bot4cj.math
|
||||
import std.io.*
|
||||
import std.fs.*
|
||||
import std.math.*
|
||||
import std.time.*
|
||||
import std.random.*
|
||||
import std.collection.*
|
||||
|
||||
/* -------- 准备与门的测试用例 --------
|
||||
* INFO(20260524-qing): 输入{2},输出{1} */
|
||||
public func prepare_andDoor(m: Matrix1, t: Matrix1) {
|
||||
let random = Random(UInt64(DateTime.now().nanosecond))
|
||||
let a = random.nextInt8(2)
|
||||
let b = random.nextInt8(2)
|
||||
let t0 = if ((a == 0 && b == 0) || (a ==1 && b == 1)) { 1.0f32 } else { 0.0f32 }
|
||||
//let t1 = if (a != b) { 1.0f32 } else { 0.0f32 }
|
||||
m.set(0, 0, FP32(Float32(a)))
|
||||
m.set(0, 1, FP32(Float32(b)))
|
||||
t.set(0, 0, FP32(t0))
|
||||
//t.set(0, 1, FP32(t1))
|
||||
}
|
||||
|
||||
/* -------- 新版神经网络推理,使用原生矩阵 -------- */
|
||||
public func forward1(nns: ArrayList<Nn1>, m: Matrix1) {
|
||||
var out_dof = Matrix1(1, 1) /* ---- 给一个值占位 ----*/
|
||||
for (idx in 0..nns.size) {
|
||||
let x = if (idx == 0) { m } else { out_dof }
|
||||
out_dof= nns[idx].forward(x)
|
||||
}
|
||||
return out_dof
|
||||
}
|
||||
|
||||
/* 【神经网络训练】
|
||||
* description: 输入一组原始数据和教师数据,进行神经网络训练
|
||||
* @param nns: 输入神经网络结构
|
||||
* @param m: 输入一组原始数据
|
||||
* @param t: 输入一组教师数据
|
||||
* @param discount: 输入本次训练的折扣系数
|
||||
*/
|
||||
public func train1(nns: ArrayList<Nn1>, m: Matrix1, t: Matrix1, discount: Float32) {
|
||||
|
||||
/* ======== 这个对象用于暂时占位 ======== */
|
||||
var res = Matrix1(1, 1)
|
||||
|
||||
/* ======== 正向传播 ======== */
|
||||
for (idx in 0..nns.size) {
|
||||
let x = if (idx == 0) { m } else { res }
|
||||
res = nns[idx].forward(x)
|
||||
}
|
||||
|
||||
/* ======== 计算误差 ======== */
|
||||
var errs = t - res
|
||||
|
||||
/* ======== 反向传播 ======== */
|
||||
for (idx in (nns.size-1)..=0:-1) {
|
||||
errs = nns[idx].backward(errs)
|
||||
}
|
||||
|
||||
/* ======== 梯度更新 ======== */
|
||||
for (idx in 0..nns.size) {
|
||||
nns[idx].update(discount)
|
||||
}
|
||||
|
||||
/* ======== 计算MSE损失 ========
|
||||
FIXME*(20260517-qing): 对于多个结果计算误差(目前只计算了第一个结果) */
|
||||
|
||||
let t0 = t.get(0, 0).getOrThrow().getValue()
|
||||
//let t1 = t.get(0, 1).getOrThrow().getValue()
|
||||
let r0 = res.get(0, 0).getOrThrow().getValue()
|
||||
//let r1 = res.get(0, 1).getOrThrow().getValue()
|
||||
let loss0 = 0.5 * (t0 - r0) * (t0 - r0)
|
||||
//let loss1 = 0.5 * (t1 - r1) * (t1 - r1)
|
||||
//return (loss0 + loss1) / 2.0
|
||||
return loss0
|
||||
}
|
||||
|
||||
/* -------- (旧版)加载多层神经网络 --------
|
||||
TODO: 也许用类封装
|
||||
INFO(20260517-qing): 新版使用仓颉原生矩阵,然后与MatrixObject类互转 */
|
||||
//public func load_nns(filename: String): ArrayList<Nn> {
|
||||
// let nns = ArrayList<Nn>()
|
||||
// let file = File(filename, OpenMode.Read)
|
||||
// let sr = StringReader<InputStream>(file)
|
||||
// while (true) {
|
||||
// if (let Some(S) <- sr.readUntil(r' ')) {
|
||||
// var S2 = S.trimAscii()
|
||||
// if (S2 == "FullConnectLayer") {
|
||||
// nns.add(Nn.load(sr)) }
|
||||
// else if (S2 == "") { break }
|
||||
// else { throw IllegalArgumentException() }
|
||||
// }
|
||||
// else { break }
|
||||
// }
|
||||
|
||||
// file.close()
|
||||
// return nns
|
||||
//}
|
||||
|
||||
/* -------- (旧版)储存多层神经网络 -------- */
|
||||
//public func save_nns(nns: ArrayList<Nn>, filename: String) {
|
||||
//
|
||||
// let file: File
|
||||
// if (exists(filename)) {
|
||||
// file = File(filename, OpenMode.Write) }
|
||||
// else { file= File.create(filename) }
|
||||
// for (nn in nns) { /* 存储全部的层 */
|
||||
// nn.save(file) }
|
||||
// file.close()
|
||||
//}
|
||||
|
||||
/* -------- 旧版神经网络推理 -------- */
|
||||
//public func forward(nns: ArrayList<Nn>, m: Matrix) {
|
||||
// var out_dof = Matrix(1, 1)
|
||||
// for (idx in 0..nns.size) {
|
||||
// let x = if (idx == 0) { m } else { out_dof }
|
||||
// let res = nns[idx].forward(x)
|
||||
// out_dof = nns[idx].fdo(res)
|
||||
// }
|
||||
// return out_dof
|
||||
//}
|
||||
|
||||
/* -------- 旧版的神经网络训练 -------- */
|
||||
//public func train(nns: ArrayList<Nn>, m: Matrix, t: Matrix, discount: Float32) {
|
||||
//
|
||||
// /* ----------------- */
|
||||
// let outs_dof = ArrayList<Matrix> ()
|
||||
// let outs_dotf = ArrayList<Matrix> ()
|
||||
//
|
||||
// /* -------- 正向传播 -------- */
|
||||
// for (idx in 0..nns.size) {
|
||||
// let x = if (idx == 0) { m } else { outs_dof[idx-1] }
|
||||
// let res = nns[idx].forward(x)
|
||||
// let out_dof = nns[idx].fdo(res)
|
||||
// let out_dotf = nns[idx].fdot(out_dof)
|
||||
// outs_dof.add(out_dof)
|
||||
// outs_dotf.add(out_dotf)
|
||||
// }
|
||||
//
|
||||
// /* -------- 计算误差 -------- */
|
||||
// var errs = outs_dof[nns.size-1] - t
|
||||
//
|
||||
// /* -------- 反向传播 -------- */
|
||||
// for (idx in (nns.size-1)..=0:-1) {
|
||||
// let up = if (idx == 0) { m} else { outs_dof[idx-1] }
|
||||
// let bias_grad = errs
|
||||
// //let delta = nns[idx].bak_bsgrad(errs, outs_dotf[idx])
|
||||
// let weights_grad = nns[idx].bak_wsgrad(up, bias_grad)
|
||||
// if (idx > 0) {
|
||||
// errs = nns[idx].bak_uperrs(bias_grad, outs_dotf[idx-1]) }
|
||||
// nns[idx].update(weights_grad, bias_grad, discount)
|
||||
// }
|
||||
//
|
||||
// /* -------- 输出损失 -------- */
|
||||
// unsafe {
|
||||
// let t0 = t.get(0, 0)
|
||||
// let t1 = t.get(0, 1)
|
||||
// let r0 = outs_dof[nns.size-1].get(0, 0)
|
||||
// let r1 = outs_dof[nns.size-1].get(0, 1)
|
||||
// let loss0 = -t0 * log(r0) - (1.0f32-t0) * log(1.0f32-r0);
|
||||
// let loss1 = -t1 * log(r1) - (1.0f32-t1) * log(1.0f32-r1);
|
||||
// return 0.5 * (loss0 * loss0 + loss1 * loss1)
|
||||
// }
|
||||
//
|
||||
//}
|
||||
@@ -1,418 +0,0 @@
|
||||
// pooling_layer.cj
|
||||
package bot4cj.math
|
||||
|
||||
import std.math.*
|
||||
import std.collection.*
|
||||
|
||||
/**
|
||||
* 池化层类
|
||||
* 设计风格与 Nn1 和 ConvLayer 保持一致
|
||||
*/
|
||||
public class PoolingLayer {
|
||||
|
||||
// ==================== 层配置参数 ====================
|
||||
let nChannels: Int64
|
||||
let poolHeight: Int64
|
||||
let poolWidth: Int64
|
||||
let stride: Int64
|
||||
let padding: Int64
|
||||
let poolType: String
|
||||
let f: Af
|
||||
let learningRateWeight: FP32
|
||||
let learningRateBias: FP32
|
||||
|
||||
// ==================== 可学习参数 ====================
|
||||
var weights: Matrix1
|
||||
var biases: Matrix1
|
||||
|
||||
// ==================== 运行时缓存 ====================
|
||||
var prevInput: Option<Matrix1> = Option.None
|
||||
var prevInputShape: Option<(Int64, Int64, Int64, Int64)> = Option.None
|
||||
var prevCombination: Option<Matrix1> = Option.None
|
||||
var prevOutput: Option<Matrix1> = Option.None
|
||||
var prevOutputShape: Option<(Int64, Int64, Int64, Int64)> = Option.None
|
||||
var prevMaxIndices: Option<Matrix1> = Option.None
|
||||
|
||||
var weightsGrad: Option<Matrix1> = Option.None
|
||||
var biasGrad: Option<Matrix1> = Option.None
|
||||
|
||||
// ==================== 构造函数 ====================
|
||||
public init(
|
||||
nChannels: Int64,
|
||||
poolSize: (Int64, Int64),
|
||||
stride: Int64,
|
||||
padding: Int64,
|
||||
poolType: String,
|
||||
activation: Af,
|
||||
lrWeight: FP32,
|
||||
lrBias: FP32
|
||||
) {
|
||||
this.nChannels = nChannels
|
||||
poolHeight = poolSize[0]
|
||||
poolWidth = poolSize[1]
|
||||
this.stride = stride
|
||||
this.padding = padding
|
||||
this.poolType = poolType
|
||||
this.f = activation
|
||||
learningRateWeight = lrWeight
|
||||
learningRateBias = lrBias
|
||||
|
||||
// 初始化权值
|
||||
weights = Matrix1(1, nChannels)
|
||||
if (poolType == "average") {
|
||||
let avgWeight = FP32(1.0) / FP32(Float32(poolHeight * poolWidth))
|
||||
for (c in 0..nChannels) {
|
||||
weights.set(0, c, avgWeight)
|
||||
}
|
||||
} else {
|
||||
for (c in 0..nChannels) {
|
||||
weights.set(0, c, FP32(1.0))
|
||||
}
|
||||
}
|
||||
|
||||
// 初始化偏置为 0
|
||||
biases = Matrix1(1, nChannels)
|
||||
for (c in 0..nChannels) {
|
||||
biases.set(0, c, FP32(0.0))
|
||||
}
|
||||
}
|
||||
|
||||
public init(
|
||||
nChannels: Int64,
|
||||
poolSize: (Int64, Int64),
|
||||
poolType: String,
|
||||
activation: Af
|
||||
) {
|
||||
this.nChannels = nChannels
|
||||
poolHeight = poolSize[0]
|
||||
poolWidth = poolSize[1]
|
||||
stride = poolHeight
|
||||
padding = 0
|
||||
this.poolType = poolType
|
||||
this.f = activation
|
||||
learningRateWeight = FP32(0.1)
|
||||
learningRateBias = FP32(0.1)
|
||||
|
||||
weights = Matrix1(1, nChannels)
|
||||
if (poolType == "average") {
|
||||
let avgWeight = FP32(1.0) / FP32(Float32(poolHeight * poolWidth))
|
||||
for (c in 0..nChannels) {
|
||||
weights.set(0, c, avgWeight)
|
||||
}
|
||||
} else {
|
||||
for (c in 0..nChannels) {
|
||||
weights.set(0, c, FP32(1.0))
|
||||
}
|
||||
}
|
||||
|
||||
biases = Matrix1(1, nChannels)
|
||||
for (c in 0..nChannels) {
|
||||
biases.set(0, c, FP32(0.0))
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== 前向传播 ====================
|
||||
public func forward(x: Matrix1, inputShape: (Int64, Int64, Int64, Int64)): Matrix1 {
|
||||
prevInput = Option.Some(x)
|
||||
prevInputShape = Option.Some(inputShape)
|
||||
|
||||
let batch = inputShape[0]
|
||||
let inH = inputShape[1]
|
||||
let inW = inputShape[2]
|
||||
|
||||
let outH = (inH + 2 * padding - poolHeight) / stride + 1
|
||||
let outW = (inW + 2 * padding - poolWidth) / stride + 1
|
||||
let outSize = batch * outH * outW
|
||||
|
||||
var combination = Matrix1(outSize, nChannels)
|
||||
var output = Matrix1(outSize, nChannels)
|
||||
var maxIndices = Matrix1(outSize, nChannels)
|
||||
|
||||
for (b in 0..batch) {
|
||||
for (oh in 0..outH) {
|
||||
for (ow in 0..outW) {
|
||||
let outRow = b * outH * outW + oh * outW + ow
|
||||
|
||||
for (ch in 0..nChannels) {
|
||||
var pooledVal = FP32(0.0)
|
||||
var maxIdx = 0
|
||||
|
||||
if (poolType == "max") {
|
||||
var maxVal = FP32(-1.0e30)
|
||||
var kernelIdx = 0
|
||||
for (kh in 0..poolHeight) {
|
||||
for (kw in 0..poolWidth) {
|
||||
let ih = oh * stride + kh - padding
|
||||
let iw = ow * stride + kw - padding
|
||||
|
||||
if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
|
||||
let inRow = b * inH * inW + ih * inW + iw
|
||||
let val = x.get(inRow, ch).getOrThrow()
|
||||
if (val.getValue() > maxVal.getValue()) {
|
||||
maxVal = val
|
||||
maxIdx = kernelIdx
|
||||
}
|
||||
}
|
||||
kernelIdx += 1
|
||||
}
|
||||
}
|
||||
pooledVal = maxVal
|
||||
} else {
|
||||
var sum = FP32(0.0)
|
||||
var count = 0
|
||||
for (kh in 0..poolHeight) {
|
||||
for (kw in 0..poolWidth) {
|
||||
let ih = oh * stride + kh - padding
|
||||
let iw = ow * stride + kw - padding
|
||||
|
||||
if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
|
||||
let inRow = b * inH * inW + ih * inW + iw
|
||||
sum = sum + x.get(inRow, ch).getOrThrow()
|
||||
count += 1
|
||||
}
|
||||
}
|
||||
}
|
||||
if (count > 0) {
|
||||
pooledVal = sum / FP32(Float32(count))
|
||||
}
|
||||
}
|
||||
|
||||
combination.set(outRow, ch, pooledVal)
|
||||
maxIndices.set(outRow, ch, FP32(Float32(maxIdx)))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 逐通道应用权值和偏置
|
||||
for (r in 0..outSize) {
|
||||
for (ch in 0..nChannels) {
|
||||
let val = combination.get(r, ch).getOrThrow()
|
||||
let w = weights.get(0, ch).getOrThrow()
|
||||
let b = biases.get(0, ch).getOrThrow()
|
||||
combination.set(r, ch, val * w + b)
|
||||
}
|
||||
}
|
||||
|
||||
prevCombination = Option.Some(combination)
|
||||
|
||||
output = f.forward(combination)
|
||||
prevOutput = Option.Some(output)
|
||||
prevOutputShape = Option.Some((batch, outH, outW, nChannels))
|
||||
prevMaxIndices = Option.Some(maxIndices)
|
||||
|
||||
return output
|
||||
}
|
||||
|
||||
// ==================== 反向传播 ====================
|
||||
public func backward(errs: Matrix1): Matrix1 {
|
||||
var x = Matrix1(1, 1)
|
||||
match (prevInput) {
|
||||
case Some(val) => x = val
|
||||
case _ => throw Exception("No input data")
|
||||
}
|
||||
|
||||
var shape = (0, 0, 0, 0)
|
||||
match (prevInputShape) {
|
||||
case Some(val) => shape = val
|
||||
case _ => throw Exception("No input shape")
|
||||
}
|
||||
|
||||
var combination = Matrix1(1, 1)
|
||||
match (prevCombination) {
|
||||
case Some(val) => combination = val
|
||||
case _ => throw Exception("No combination")
|
||||
}
|
||||
|
||||
var outShape = (0, 0, 0, 0)
|
||||
match (prevOutputShape) {
|
||||
case Some(val) => outShape = val
|
||||
case _ => throw Exception("No output shape")
|
||||
}
|
||||
|
||||
var maxIndices = Matrix1(1, 1)
|
||||
let hasMaxIndices = match (prevMaxIndices) {
|
||||
case Some(_) => true
|
||||
case _ => false
|
||||
}
|
||||
if (hasMaxIndices) {
|
||||
maxIndices = match (prevMaxIndices) {
|
||||
case Some(val) => val
|
||||
case _ => Matrix1(1, 1)
|
||||
}
|
||||
}
|
||||
|
||||
let batch = shape[0]
|
||||
let inH = shape[1]
|
||||
let inW = shape[2]
|
||||
let outH = outShape[1]
|
||||
let outW = outShape[2]
|
||||
|
||||
let delta = errs * f.backward(combination)
|
||||
|
||||
calcWeightBiasGradient(delta, combination)
|
||||
|
||||
let errsUp = calcUpsampledError(delta, maxIndices, batch, inH, inW, outH, outW)
|
||||
|
||||
return errsUp
|
||||
}
|
||||
|
||||
/**
|
||||
* 计算权值和偏置梯度
|
||||
*/
|
||||
private func calcWeightBiasGradient(delta: Matrix1, combination: Matrix1): Unit {
|
||||
let outSize = delta.get_row()
|
||||
|
||||
var wGrad = Matrix1(1, nChannels)
|
||||
var bGrad = Matrix1(1, nChannels)
|
||||
|
||||
for (ch in 0..nChannels) {
|
||||
var wSum = FP32(0.0)
|
||||
var bSum = FP32(0.0)
|
||||
for (r in 0..outSize) {
|
||||
let d = delta.get(r, ch).getOrThrow()
|
||||
bSum = bSum + d
|
||||
let comb = combination.get(r, ch).getOrThrow()
|
||||
wSum = wSum + d * comb
|
||||
}
|
||||
wGrad.set(0, ch, wSum)
|
||||
bGrad.set(0, ch, bSum)
|
||||
}
|
||||
|
||||
weightsGrad = Option.Some(wGrad)
|
||||
biasGrad = Option.Some(bGrad)
|
||||
}
|
||||
|
||||
/**
|
||||
* 上采样误差(池化层的反向传播核心)
|
||||
*/
|
||||
private func calcUpsampledError(
|
||||
delta: Matrix1,
|
||||
maxIndices: Matrix1,
|
||||
batch: Int64,
|
||||
inH: Int64,
|
||||
inW: Int64,
|
||||
outH: Int64,
|
||||
outW: Int64
|
||||
): Matrix1 {
|
||||
let inSize = batch * inH * inW
|
||||
var errsUp = Matrix1(inSize, nChannels)
|
||||
|
||||
for (b in 0..batch) {
|
||||
for (oh in 0..outH) {
|
||||
for (ow in 0..outW) {
|
||||
let deltaRow = b * outH * outW + oh * outW + ow
|
||||
|
||||
for (ch in 0..nChannels) {
|
||||
let deltaVal = delta.get(deltaRow, ch).getOrThrow()
|
||||
let w = weights.get(0, ch).getOrThrow()
|
||||
|
||||
if (poolType == "max") {
|
||||
let maxIdx = Int64(maxIndices.get(deltaRow, ch).getOrThrow().getValue())
|
||||
|
||||
var kernelIdx = 0
|
||||
var found = false
|
||||
for (kh in 0..poolHeight) {
|
||||
for (kw in 0..poolWidth) {
|
||||
if (kernelIdx == maxIdx && !found) {
|
||||
let ih = oh * stride + kh - padding
|
||||
let iw = ow * stride + kw - padding
|
||||
|
||||
if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
|
||||
let inRow = b * inH * inW + ih * inW + iw
|
||||
errsUp.set(inRow, ch, deltaVal * w)
|
||||
found = true
|
||||
}
|
||||
}
|
||||
kernelIdx += 1
|
||||
}
|
||||
}
|
||||
} else {
|
||||
var count = 0
|
||||
for (kh in 0..poolHeight) {
|
||||
for (kw in 0..poolWidth) {
|
||||
let ih = oh * stride + kh - padding
|
||||
let iw = ow * stride + kw - padding
|
||||
if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
|
||||
count += 1
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (count > 0) {
|
||||
let avgDelta = deltaVal * w / FP32(Float32(count))
|
||||
for (kh in 0..poolHeight) {
|
||||
for (kw in 0..poolWidth) {
|
||||
let ih = oh * stride + kh - padding
|
||||
let iw = ow * stride + kw - padding
|
||||
if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
|
||||
let inRow = b * inH * inW + ih * inW + iw
|
||||
let oldVal = errsUp.get(inRow, ch).getOrThrow()
|
||||
errsUp.set(inRow, ch, oldVal + avgDelta)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return errsUp
|
||||
}
|
||||
|
||||
// ==================== 参数更新 ====================
|
||||
public func update(discount: Float32): Unit {
|
||||
var wGrad = Matrix1(1, 1)
|
||||
match (weightsGrad) {
|
||||
case Some(val) => wGrad = val
|
||||
case _ => throw Exception("No weight gradients")
|
||||
}
|
||||
|
||||
var bGrad = Matrix1(1, 1)
|
||||
match (biasGrad) {
|
||||
case Some(val) => bGrad = val
|
||||
case _ => throw Exception("No bias gradients")
|
||||
}
|
||||
|
||||
for (ch in 0..nChannels) {
|
||||
let oldW = weights.get(0, ch).getOrThrow()
|
||||
let gW = wGrad.get(0, ch).getOrThrow()
|
||||
weights.set(0, ch, oldW + gW * learningRateWeight * FP32(discount))
|
||||
}
|
||||
|
||||
for (ch in 0..nChannels) {
|
||||
let oldB = biases.get(0, ch).getOrThrow()
|
||||
let gB = bGrad.get(0, ch).getOrThrow()
|
||||
biases.set(0, ch, oldB + gB * learningRateBias * FP32(discount))
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== 形状访问器 ====================
|
||||
public func getInShape(): (Int64, Int64) {
|
||||
(nChannels, poolHeight * poolWidth)
|
||||
}
|
||||
|
||||
public func getOutShape(): Int64 {
|
||||
nChannels
|
||||
}
|
||||
|
||||
public func getPoolType(): String {
|
||||
poolType
|
||||
}
|
||||
|
||||
public func getLR(): FP32 {
|
||||
learningRateWeight
|
||||
}
|
||||
|
||||
public func getFType(): String {
|
||||
f.get_type()
|
||||
}
|
||||
|
||||
public func calculateOutputSize(inputHeight: Int64, inputWidth: Int64): (Int64, Int64) {
|
||||
let outH = (inputHeight + 2 * padding - poolHeight) / stride + 1
|
||||
let outW = (inputWidth + 2 * padding - poolWidth) / stride + 1
|
||||
(outH, outW)
|
||||
}
|
||||
}
|
||||
@@ -1,310 +0,0 @@
|
||||
// test_cnn.cj - 使用 Adam 优化器的版本(所有层)
|
||||
package bot4cj
|
||||
import bot4cj.math.*
|
||||
import std.math.*
|
||||
import std.random.*
|
||||
import std.collection.*
|
||||
|
||||
func generateSolidColorData(numSamples: Int64): (Array<Matrix1>, Array<Int64>) {
|
||||
var inputs = Array(numSamples, { _ => Matrix1(1, 1) })
|
||||
var labels = Array(numSamples, { _ => 0 })
|
||||
let rng = Random()
|
||||
|
||||
for (i in 0..numSamples) {
|
||||
let colorClass = Int64(rng.nextFloat64() * 3.0)
|
||||
labels[i] = colorClass
|
||||
|
||||
let h = 8
|
||||
let w = 8
|
||||
var matrix = Matrix1(h * w, 1)
|
||||
|
||||
let colorValue = match (colorClass) {
|
||||
case 0 => FP32(1.0)
|
||||
case 1 => FP32(0.5)
|
||||
case _ => FP32(0.0)
|
||||
}
|
||||
|
||||
for (r in 0..h * w) {
|
||||
let noise = FP32(Float32(rng.nextFloat64() - 0.5) * 0.1)
|
||||
var val = colorValue + noise
|
||||
if (val.getValue() > 1.0) {
|
||||
val = FP32(1.0)
|
||||
} else if (val.getValue() < 0.0) {
|
||||
val = FP32(0.0)
|
||||
}
|
||||
matrix.set(r, 0, val)
|
||||
}
|
||||
|
||||
inputs[i] = matrix
|
||||
}
|
||||
|
||||
(inputs, labels)
|
||||
}
|
||||
|
||||
func printSeparator(): Unit {
|
||||
println("==================================================")
|
||||
}
|
||||
|
||||
func testCNN(): Unit {
|
||||
printSeparator()
|
||||
println("CNN 纯色矩阵识别测试 (所有层使用 Adam 优化器)")
|
||||
printSeparator()
|
||||
println()
|
||||
|
||||
let nTrainSamples = 200
|
||||
let nTestSamples = 60
|
||||
let nEpochs = 50
|
||||
let batchSize = 10
|
||||
|
||||
println("训练样本数: ${nTrainSamples}")
|
||||
println("测试样本数: ${nTestSamples}")
|
||||
println("训练轮数: ${nEpochs}")
|
||||
println("Batch 大小: ${batchSize}")
|
||||
println()
|
||||
|
||||
println("创建网络结构...")
|
||||
println("Input(8x8x1) -> Conv(3x3,1->4) -> Pool(2x2) -> Conv(3x3,4->8) -> Flatten -> FC(8,3)")
|
||||
println()
|
||||
|
||||
let lr = FP32(0.001)
|
||||
|
||||
var conv1 = ConvLayer(1, 4, (3, 3), 1, 0, lr, Af.Leaky_ReLU(FP32(0.01)), Option.None)
|
||||
var pool1 = PoolingLayer(4, (2, 2), 2, 0, "max", Af.ReLU, lr, lr)
|
||||
var conv2 = ConvLayer(4, 8, (3, 3), 1, 0, lr, Af.Leaky_ReLU(FP32(0.01)), Option.None)
|
||||
var flatten = FlattenLayer(8, 1, 1)
|
||||
var fc = Nn1(8, 3, Float32(0.4), Af.Softmax)
|
||||
|
||||
// 创建 Adam 优化器(为每个可训练层创建独立的优化器)
|
||||
var adamConv1 = Adam(FP32(0.001)) // 卷积层1
|
||||
var adamConv2 = Adam(FP32(0.001)) // 卷积层2
|
||||
var adamFC = Adam(FP32(0.0005)) // 全连接层(使用更小的学习率)
|
||||
|
||||
println("网络创建完成 (Conv1, Conv2, FC 使用 Adam 优化器)")
|
||||
println()
|
||||
|
||||
println("生成数据...")
|
||||
let (trainInputs, trainLabels) = generateSolidColorData(nTrainSamples)
|
||||
let (testInputs, testLabels) = generateSolidColorData(nTestSamples)
|
||||
println("训练集: ${nTrainSamples} 样本")
|
||||
println("测试集: ${nTestSamples} 样本")
|
||||
println()
|
||||
|
||||
println("开始训练...")
|
||||
println()
|
||||
|
||||
for (epoch in 0..nEpochs) {
|
||||
var epochLoss = Float32(0.0)
|
||||
var epochCorrect = 0
|
||||
|
||||
for (i in 0..nTrainSamples) {
|
||||
let x = trainInputs[i]
|
||||
let label = trainLabels[i]
|
||||
|
||||
// ===== 前向传播 =====
|
||||
let c1Out = conv1.forward(x, (1, 8, 8, 1))
|
||||
let p1Out = pool1.forward(c1Out, (1, 6, 6, 4))
|
||||
let c2Out = conv2.forward(p1Out, (1, 3, 3, 4))
|
||||
let flatOut = flatten.forward(c2Out, (1, 1, 1, 8))
|
||||
|
||||
// 调试:确认 flatten.forward 执行成功
|
||||
if (i == 0 && epoch == 0) {
|
||||
println("DEBUG: flatten.forward done, flatOut rows: ${flatOut.get_row()}, cols: ${flatOut.get_col()}")
|
||||
}
|
||||
|
||||
var fcInput = Matrix1(1, 8)
|
||||
for (j in 0..8) {
|
||||
fcInput.set(0, j, flatOut.get(j, 0).getOrThrow())
|
||||
}
|
||||
let fcOut = fc.forward(fcInput)
|
||||
|
||||
// Softmax
|
||||
var softmax = Matrix1(1, 3)
|
||||
var maxVal = Float32(-1.0e30)
|
||||
for (c in 0..3) {
|
||||
let raw = fcOut.get(0, c).getOrThrow().getValue()
|
||||
if (raw > maxVal) {
|
||||
maxVal = raw
|
||||
}
|
||||
}
|
||||
var expSum = Float32(0.0)
|
||||
for (c in 0..3) {
|
||||
let raw = fcOut.get(0, c).getOrThrow().getValue()
|
||||
let ev = Float32(exp(raw - maxVal))
|
||||
softmax.set(0, c, FP32(ev))
|
||||
expSum = expSum + ev
|
||||
}
|
||||
for (c in 0..3) {
|
||||
let v = softmax.get(0, c).getOrThrow().getValue()
|
||||
softmax.set(0, c, FP32(v / expSum))
|
||||
}
|
||||
|
||||
let probVal = softmax.get(0, label).getOrThrow().getValue()
|
||||
let loss = -log(max(probVal, Float32(1e-10)))
|
||||
epochLoss = epochLoss + loss
|
||||
|
||||
var maxProb = Float32(-1.0)
|
||||
var predClass = 0
|
||||
for (c in 0..3) {
|
||||
let p = softmax.get(0, c).getOrThrow().getValue()
|
||||
if (p > maxProb) {
|
||||
maxProb = p
|
||||
predClass = Int64(c)
|
||||
}
|
||||
}
|
||||
if (predClass == label) {
|
||||
epochCorrect += 1
|
||||
}
|
||||
|
||||
// ===== 反向传播 =====
|
||||
var delta = Matrix1(1, 3)
|
||||
for (c in 0..3) {
|
||||
let prob = softmax.get(0, c).getOrThrow()
|
||||
if (Int64(c) == label) {
|
||||
delta.set(0, c, prob - FP32(1.0))
|
||||
} else {
|
||||
delta.set(0, c, prob)
|
||||
}
|
||||
}
|
||||
|
||||
let fcDelta = fc.backward(delta)
|
||||
|
||||
// 调试:确认 fcDelta 形状
|
||||
if (i == 0 && epoch == 0) {
|
||||
println("DEBUG: fc.backward done, fcDelta rows: ${fcDelta.get_row()}, cols: ${fcDelta.get_col()}")
|
||||
}
|
||||
|
||||
var flatDelta = Matrix1(8, 1)
|
||||
for (j in 0..8) {
|
||||
flatDelta.set(j, 0, fcDelta.get(0, j).getOrThrow())
|
||||
}
|
||||
|
||||
// 调试:确认 flatDelta 形状
|
||||
if (i == 0 && epoch == 0) {
|
||||
println("DEBUG: flatDelta rows: ${flatDelta.get_row()}, cols: ${flatDelta.get_col()}")
|
||||
}
|
||||
|
||||
let c2Delta = flatten.backward(flatDelta)
|
||||
let p1Delta = conv2.backward(c2Delta)
|
||||
let c1Delta = pool1.backward(p1Delta)
|
||||
let _ = conv1.backward(c1Delta)
|
||||
|
||||
// ===== 参数更新 =====
|
||||
if ((i + 1) % batchSize == 0) {
|
||||
// 所有可训练层使用 Adam 更新
|
||||
adamConv1.updateConvLayer(conv1) // 卷积层1 - Adam
|
||||
adamConv2.updateConvLayer(conv2) // 卷积层2 - Adam
|
||||
adamFC.updateFCLayer(fc) // 全连接层 - Adam
|
||||
|
||||
// 池化层使用 SGD 更新(如果需要的话)
|
||||
pool1.update(Float32(1.0))
|
||||
}
|
||||
}
|
||||
|
||||
// Epoch 结束时的更新
|
||||
adamConv1.updateConvLayer(conv1)
|
||||
adamConv2.updateConvLayer(conv2)
|
||||
adamFC.updateFCLayer(fc)
|
||||
pool1.update(Float32(1.0))
|
||||
|
||||
let avgLoss = epochLoss / Float32(nTrainSamples)
|
||||
let accuracy = Float32(epochCorrect) / Float32(nTrainSamples) * Float32(100.0)
|
||||
|
||||
println("Epoch ${epoch + 1}/${nEpochs} - Loss: ${avgLoss}, Accuracy: ${accuracy}%")
|
||||
|
||||
// 打印 Adam 的步数信息
|
||||
if (epoch == 0) {
|
||||
println(" Adam Step: ${adamConv1.getStep()}")
|
||||
}
|
||||
}
|
||||
|
||||
println("")
|
||||
println("开始测试...")
|
||||
println("")
|
||||
|
||||
var testCorrect = 0
|
||||
var confusion = Array(3, { _ => Array(3, { _ => 0 }) })
|
||||
|
||||
for (i in 0..nTestSamples) {
|
||||
let x = testInputs[i]
|
||||
let label = testLabels[i]
|
||||
|
||||
let c1Out = conv1.forward(x, (1, 8, 8, 1))
|
||||
let p1Out = pool1.forward(c1Out, (1, 6, 6, 4))
|
||||
let c2Out = conv2.forward(p1Out, (1, 3, 3, 4))
|
||||
let flatOut = flatten.forward(c2Out, (1, 1, 1, 8))
|
||||
|
||||
var fcInput = Matrix1(1, 8)
|
||||
for (j in 0..8) {
|
||||
fcInput.set(0, j, flatOut.get(j, 0).getOrThrow())
|
||||
}
|
||||
let fcOut = fc.forward(fcInput)
|
||||
|
||||
var softmax = Matrix1(1, 3)
|
||||
var maxVal = Float32(-1.0e30)
|
||||
for (c in 0..3) {
|
||||
let raw = fcOut.get(0, c).getOrThrow().getValue()
|
||||
if (raw > maxVal) {
|
||||
maxVal = raw
|
||||
}
|
||||
}
|
||||
var expSum = Float32(0.0)
|
||||
for (c in 0..3) {
|
||||
let raw = fcOut.get(0, c).getOrThrow().getValue()
|
||||
let ev = Float32(exp(raw - maxVal))
|
||||
softmax.set(0, c, FP32(ev))
|
||||
expSum = expSum + ev
|
||||
}
|
||||
for (c in 0..3) {
|
||||
let v = softmax.get(0, c).getOrThrow().getValue()
|
||||
softmax.set(0, c, FP32(v / expSum))
|
||||
}
|
||||
|
||||
var maxProb = Float32(-1.0)
|
||||
var predClass = 0
|
||||
for (c in 0..3) {
|
||||
let p = softmax.get(0, c).getOrThrow().getValue()
|
||||
if (p > maxProb) {
|
||||
maxProb = p
|
||||
predClass = Int64(c)
|
||||
}
|
||||
}
|
||||
|
||||
confusion[label][predClass] += 1
|
||||
if (predClass == label) {
|
||||
testCorrect += 1
|
||||
}
|
||||
}
|
||||
|
||||
let testAcc = Float32(testCorrect) / Float32(nTestSamples) * Float32(100.0)
|
||||
println("测试准确率: ${testAcc}% (${testCorrect}/${nTestSamples})")
|
||||
print("")
|
||||
|
||||
let classNames = ["红色(高)", "绿色(中)", "蓝色(低)"]
|
||||
println("各类别准确率:")
|
||||
for (c in 0..3) {
|
||||
var total = 0
|
||||
for (p in 0..3) {
|
||||
total += confusion[c][p]
|
||||
}
|
||||
if (total > 0) {
|
||||
let acc = Float32(confusion[c][c]) / Float32(total) * Float32(100.0)
|
||||
print(" ${classNames[c]}: ${confusion[c][c]}/${total} (${acc}%)")
|
||||
}
|
||||
}
|
||||
println("")
|
||||
|
||||
println("混淆矩阵:")
|
||||
println(" 预测红 预测绿 预测蓝")
|
||||
for (t in 0..3) {
|
||||
var line = "真实${classNames[t]}:"
|
||||
for (p in 0..3) {
|
||||
line = line + " ${confusion[t][p]} "
|
||||
}
|
||||
print(line)
|
||||
}
|
||||
println("")
|
||||
|
||||
printSeparator()
|
||||
println("测试完成!")
|
||||
printSeparator()
|
||||
}
|
||||
@@ -1,517 +0,0 @@
|
||||
// test_geometry_cnn.cj - 几何图形分类测试(Adam 优化器版本)
|
||||
package bot4cj
|
||||
import bot4cj.math.*
|
||||
import std.math.*
|
||||
import std.random.*
|
||||
import std.collection.*
|
||||
|
||||
/**
|
||||
* 计算圆形像素值
|
||||
*/
|
||||
func calcCirclePixel(x: Int64, y: Int64, center: Float64, radius: Float64, offsetX: Float64, offsetY: Float64): FP32 {
|
||||
let dx = Float64(x) - center + offsetX
|
||||
let dy = Float64(y) - center + offsetY
|
||||
let dist = sqrt(dx * dx + dy * dy)
|
||||
|
||||
if (dist <= radius) {
|
||||
let intensity = 1.0 - (dist / radius) * 0.3
|
||||
FP32(Float32(intensity))
|
||||
} else if (dist <= radius + 1.5) {
|
||||
let edgeDist = dist - radius
|
||||
let intensity = max(0.0, 1.0 - edgeDist / 1.5) * 0.7
|
||||
FP32(Float32(intensity))
|
||||
} else {
|
||||
FP32(0.0)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 计算正方形像素值
|
||||
*/
|
||||
func calcSquarePixel(x: Int64, y: Int64, center: Float64, radius: Float64, offsetX: Float64, offsetY: Float64): FP32 {
|
||||
let halfSize = radius * 0.8
|
||||
let left = center - halfSize + offsetX
|
||||
let right = center + halfSize + offsetX
|
||||
let top = center - halfSize + offsetY
|
||||
let bottom = center + halfSize + offsetY
|
||||
|
||||
let fx = Float64(x)
|
||||
let fy = Float64(y)
|
||||
|
||||
if (fx >= left && fx <= right && fy >= top && fy <= bottom) {
|
||||
let distToEdge = min(
|
||||
min(fx - left, right - fx),
|
||||
min(fy - top, bottom - fy)
|
||||
)
|
||||
let intensity = min(1.0, 0.7 + distToEdge / halfSize * 0.3)
|
||||
FP32(Float32(intensity))
|
||||
} else {
|
||||
let distX = max(0.0, max(left - fx, fx - right))
|
||||
let distY = max(0.0, max(top - fy, fy - bottom))
|
||||
let edgeDist = sqrt(distX * distX + distY * distY)
|
||||
if (edgeDist <= 1.5) {
|
||||
let intensity = max(0.0, 1.0 - edgeDist / 1.5) * 0.6
|
||||
FP32(Float32(intensity))
|
||||
} else {
|
||||
FP32(0.0)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 计算三角形像素值
|
||||
*/
|
||||
func calcTrianglePixel(x: Int64, y: Int64, center: Float64, radius: Float64, offsetX: Float64, offsetY: Float64): FP32 {
|
||||
let height = radius * 2.0
|
||||
let baseWidth = height * 1.2
|
||||
let topY = center - height * 0.5 + offsetY
|
||||
let bottomY = center + height * 0.5 + offsetY
|
||||
|
||||
let fx = Float64(x)
|
||||
let fy = Float64(y)
|
||||
|
||||
if (fy >= topY && fy <= bottomY) {
|
||||
let t = (fy - topY) / (bottomY - topY)
|
||||
let halfWidthAtY = baseWidth * (1.0 - t) * 0.5
|
||||
let centerX = center + offsetX
|
||||
|
||||
if (fx >= centerX - halfWidthAtY && fx <= centerX + halfWidthAtY) {
|
||||
let distToEdge = min(
|
||||
fx - (centerX - halfWidthAtY),
|
||||
(centerX + halfWidthAtY) - fx
|
||||
)
|
||||
let intensity = min(1.0, 0.7 + distToEdge / halfWidthAtY * 0.3)
|
||||
FP32(Float32(intensity))
|
||||
} else {
|
||||
let distToLeft = (centerX - halfWidthAtY) - fx
|
||||
let distToRight = fx - (centerX + halfWidthAtY)
|
||||
let edgeDist = max(distToLeft, distToRight)
|
||||
if (edgeDist <= 1.5) {
|
||||
let intensity = max(0.0, 1.0 - edgeDist / 1.5) * 0.6
|
||||
FP32(Float32(intensity))
|
||||
} else {
|
||||
FP32(0.0)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
FP32(0.0)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 生成几何图形数据
|
||||
* 图形类型:0=圆形, 1=正方形, 2=三角形
|
||||
*/
|
||||
func generateGeometryData(numSamples: Int64): (Array<Matrix1>, Array<Int64>) {
|
||||
var inputs = Array(numSamples, { _ => Matrix1(1, 1) })
|
||||
var labels = Array(numSamples, { _ => 0 })
|
||||
let rng = Random()
|
||||
|
||||
let imgSize = 16
|
||||
|
||||
for (i in 0..numSamples) {
|
||||
let shapeType = Int64(rng.nextFloat64() * 3.0)
|
||||
labels[i] = shapeType
|
||||
|
||||
var matrix = Matrix1(imgSize * imgSize, 1)
|
||||
let center = Float64(imgSize) / 2.0
|
||||
let radius = Float64(imgSize) / 3.0 + rng.nextFloat64() * 2.0 - 1.0
|
||||
|
||||
let offsetX = (rng.nextFloat64() - 0.5) * 2.0
|
||||
let offsetY = (rng.nextFloat64() - 0.5) * 2.0
|
||||
|
||||
for (y in 0..imgSize) {
|
||||
for (x in 0..imgSize) {
|
||||
let idx = y * imgSize + x
|
||||
|
||||
var pixelValue = match (shapeType) {
|
||||
case 0 => calcCirclePixel(x, y, center, radius, offsetX, offsetY)
|
||||
case 1 => calcSquarePixel(x, y, center, radius, offsetX, offsetY)
|
||||
case _ => calcTrianglePixel(x, y, center, radius, offsetX, offsetY)
|
||||
}
|
||||
|
||||
let noise = FP32(Float32(rng.nextFloat64() - 0.5) * 0.05)
|
||||
pixelValue = pixelValue + noise
|
||||
|
||||
if (pixelValue.getValue() > 1.0) {
|
||||
pixelValue = FP32(1.0)
|
||||
} else if (pixelValue.getValue() < 0.0) {
|
||||
pixelValue = FP32(0.0)
|
||||
}
|
||||
|
||||
matrix.set(idx, 0, pixelValue)
|
||||
}
|
||||
}
|
||||
|
||||
inputs[i] = matrix
|
||||
}
|
||||
|
||||
(inputs, labels)
|
||||
}
|
||||
|
||||
/**
|
||||
* 可视化一个样本
|
||||
*/
|
||||
func showSample(matrix: Matrix1, imgSize: Int64, label: Int64): Unit {
|
||||
let shapeNames = ["圆形", "正方形", "三角形"]
|
||||
println("标签: ${shapeNames[label]}")
|
||||
|
||||
for (y in 0..imgSize) {
|
||||
var line = ""
|
||||
for (x in 0..imgSize) {
|
||||
let idx = y * imgSize + x
|
||||
let val = matrix.get(idx, 0).getOrThrow().getValue()
|
||||
if (val > 0.8) {
|
||||
line = line + "██"
|
||||
} else if (val > 0.6) {
|
||||
line = line + "▓▓"
|
||||
} else if (val > 0.4) {
|
||||
line = line + "▒▒"
|
||||
} else if (val > 0.2) {
|
||||
line = line + "░░"
|
||||
} else {
|
||||
line = line + " "
|
||||
}
|
||||
}
|
||||
println(line)
|
||||
}
|
||||
println("")
|
||||
}
|
||||
|
||||
/**
|
||||
* Softmax 辅助函数
|
||||
*/
|
||||
func computeSoftmax(input: Matrix1): Matrix1 {
|
||||
let cols = input.get_col()
|
||||
var output = Matrix1(1, cols)
|
||||
|
||||
var maxVal = Float32(-1.0e30)
|
||||
for (c in 0..cols) {
|
||||
let raw = input.get(0, c).getOrThrow().getValue()
|
||||
if (raw > maxVal) {
|
||||
maxVal = raw
|
||||
}
|
||||
}
|
||||
|
||||
var expSum = Float32(0.0)
|
||||
for (c in 0..cols) {
|
||||
let raw = input.get(0, c).getOrThrow().getValue()
|
||||
let ev = Float32(exp(raw - maxVal))
|
||||
output.set(0, c, FP32(ev))
|
||||
expSum = expSum + ev
|
||||
}
|
||||
|
||||
for (c in 0..cols) {
|
||||
let v = output.get(0, c).getOrThrow().getValue()
|
||||
output.set(0, c, FP32(v / expSum))
|
||||
}
|
||||
|
||||
output
|
||||
}
|
||||
|
||||
/**
|
||||
* 获取预测类别
|
||||
*/
|
||||
func predictClass(probs: Matrix1): Int64 {
|
||||
let cols = probs.get_col()
|
||||
var maxProb = Float32(-1.0)
|
||||
var predClass = 0
|
||||
|
||||
for (c in 0..cols) {
|
||||
let p = probs.get(0, c).getOrThrow().getValue()
|
||||
if (p > maxProb) {
|
||||
maxProb = p
|
||||
predClass = Int64(c)
|
||||
}
|
||||
}
|
||||
|
||||
predClass
|
||||
}
|
||||
|
||||
/**
|
||||
* 打印分隔线
|
||||
*/
|
||||
func printLine(): Unit {
|
||||
println("============================================================")
|
||||
}
|
||||
|
||||
/**
|
||||
* 主测试函数
|
||||
*/
|
||||
func testGeometryCNN(): Unit {
|
||||
printLine()
|
||||
println("CNN 几何图形分类测试 (Adam 优化器版本)")
|
||||
printLine()
|
||||
println("")
|
||||
|
||||
let imgSize = 16
|
||||
let nTrainSamples = 200
|
||||
let nTestSamples = 60
|
||||
let nEpochs = 50
|
||||
let batchSize = 10
|
||||
let nClasses = 3
|
||||
|
||||
println("配置信息:")
|
||||
println(" 图像大小: ${imgSize}x${imgSize}")
|
||||
println(" 训练样本数: ${nTrainSamples}")
|
||||
println(" 测试样本数: ${nTestSamples}")
|
||||
println(" 训练轮数: ${nEpochs}")
|
||||
println(" Batch 大小: ${batchSize}")
|
||||
println(" 类别数: ${nClasses} (圆形, 正方形, 三角形)")
|
||||
println(" 优化器: Adam")
|
||||
println("")
|
||||
|
||||
println("网络结构:")
|
||||
println(" Input(16x16x1)")
|
||||
println(" -> Conv1(5x5, 1->8) -> Pool1(2x2)")
|
||||
println(" -> Conv2(3x3, 8->16) -> Pool2(2x2)")
|
||||
println(" -> Flatten -> FC(64->32) -> FC(32->3)")
|
||||
println(" -> Softmax")
|
||||
println("")
|
||||
|
||||
let lr = FP32(0.01)
|
||||
|
||||
// 构建网络
|
||||
var conv1 = ConvLayer(1, 8, (5, 5), 1, 0, lr, Af.Leaky_ReLU(FP32(0.01)), Option.None)
|
||||
var pool1 = PoolingLayer(8, (2, 2), 2, 0, "max", Af.ReLU, lr, lr)
|
||||
var conv2 = ConvLayer(8, 16, (3, 3), 1, 0, lr, Af.Leaky_ReLU(FP32(0.01)), Option.None)
|
||||
var pool2 = PoolingLayer(16, (2, 2), 2, 0, "max", Af.ReLU, lr, lr)
|
||||
var flatten = FlattenLayer(16, 2, 2)
|
||||
var fc1 = Nn1(64, 32, Float32(0.3), Af.Leaky_ReLU(FP32(0.01)))
|
||||
var fc2 = Nn1(32, 3, Float32(0.3), Af.Softmax)
|
||||
|
||||
// 创建 Adam 优化器
|
||||
var adamConv1 = Adam(FP32(0.001)) // 卷积层1
|
||||
var adamConv2 = Adam(FP32(0.001)) // 卷积层2
|
||||
var adamFC1 = Adam(FP32(0.0005)) // 全连接层1
|
||||
var adamFC2 = Adam(FP32(0.0005)) // 全连接层2
|
||||
|
||||
println("网络创建完成 (Adam优化器: Conv1, Conv2, FC1, FC2)")
|
||||
println("")
|
||||
|
||||
// 生成数据
|
||||
println("生成训练数据...")
|
||||
let (trainInputs, trainLabels) = generateGeometryData(nTrainSamples)
|
||||
println("生成测试数据...")
|
||||
let (testInputs, testLabels) = generateGeometryData(nTestSamples)
|
||||
println("数据生成完成")
|
||||
println("")
|
||||
|
||||
// 显示几个样本
|
||||
println("样本预览:")
|
||||
for (i in 0..3) {
|
||||
if (i < nTrainSamples) {
|
||||
println("样本 ${i + 1}:")
|
||||
showSample(trainInputs[i], imgSize, trainLabels[i])
|
||||
}
|
||||
}
|
||||
|
||||
println("开始训练...")
|
||||
println("")
|
||||
|
||||
var bestAccuracy = Float32(0.0)
|
||||
|
||||
// 计算卷积输出尺寸
|
||||
let conv1OutH = imgSize - 4 // 16 - 5 + 1 = 12
|
||||
let conv1OutW = imgSize - 4
|
||||
let pool1OutH = conv1OutH / 2 // 6
|
||||
let pool1OutW = conv1OutW / 2
|
||||
let conv2OutH = pool1OutH - 2 // 6 - 3 + 1 = 4
|
||||
let conv2OutW = pool1OutW - 2
|
||||
let pool2OutH = conv2OutH / 2 // 2
|
||||
let pool2OutW = conv2OutW / 2
|
||||
let flatSize = 16 * pool2OutH * pool2OutW // 64
|
||||
|
||||
for (epoch in 0..nEpochs) {
|
||||
var epochLoss = Float32(0.0)
|
||||
var epochCorrect = 0
|
||||
|
||||
// 训练
|
||||
for (i in 0..nTrainSamples) {
|
||||
let x = trainInputs[i]
|
||||
let label = trainLabels[i]
|
||||
|
||||
// 前向传播
|
||||
let c1Out = conv1.forward(x, (1, imgSize, imgSize, 1))
|
||||
let p1Out = pool1.forward(c1Out, (1, conv1OutH, conv1OutW, 8))
|
||||
let c2Out = conv2.forward(p1Out, (1, pool1OutH, pool1OutW, 8))
|
||||
let p2Out = pool2.forward(c2Out, (1, conv2OutH, conv2OutW, 16))
|
||||
|
||||
let flatOut = flatten.forward(p2Out, (1, pool2OutH, pool2OutW, 16))
|
||||
|
||||
// 转换为全连接层输入
|
||||
var fc1Input = Matrix1(1, flatSize)
|
||||
for (j in 0..flatSize) {
|
||||
fc1Input.set(0, j, flatOut.get(j, 0).getOrThrow())
|
||||
}
|
||||
|
||||
let fc1Out = fc1.forward(fc1Input)
|
||||
let fc2Out = fc2.forward(fc1Out)
|
||||
|
||||
// Softmax
|
||||
let probs = computeSoftmax(fc2Out)
|
||||
|
||||
// 计算损失和准确率
|
||||
let probVal = probs.get(0, label).getOrThrow().getValue()
|
||||
let loss = -log(max(probVal, Float32(1e-10)))
|
||||
epochLoss = epochLoss + loss
|
||||
|
||||
let predClass = predictClass(probs)
|
||||
if (predClass == label) {
|
||||
epochCorrect += 1
|
||||
}
|
||||
|
||||
// 反向传播
|
||||
var delta = Matrix1(1, nClasses)
|
||||
for (c in 0..nClasses) {
|
||||
let prob = probs.get(0, c).getOrThrow()
|
||||
if (Int64(c) == label) {
|
||||
delta.set(0, c, prob - FP32(1.0))
|
||||
} else {
|
||||
delta.set(0, c, prob)
|
||||
}
|
||||
}
|
||||
|
||||
let fc2Delta = fc2.backward(delta)
|
||||
let fc1Delta = fc1.backward(fc2Delta)
|
||||
|
||||
// 转换回特征图格式
|
||||
var flatDelta = Matrix1(flatSize, 1)
|
||||
for (j in 0..flatSize) {
|
||||
flatDelta.set(j, 0, fc1Delta.get(0, j).getOrThrow())
|
||||
}
|
||||
|
||||
let p2Delta = flatten.backward(flatDelta)
|
||||
let c2Delta = pool2.backward(p2Delta)
|
||||
let p1Delta = conv2.backward(c2Delta)
|
||||
let c1Delta = pool1.backward(p1Delta)
|
||||
let _ = conv1.backward(c1Delta)
|
||||
|
||||
// 批次更新 - 使用 Adam
|
||||
if ((i + 1) % batchSize == 0) {
|
||||
adamConv1.updateConvLayer(conv1) // Conv1 - Adam
|
||||
adamConv2.updateConvLayer(conv2) // Conv2 - Adam
|
||||
adamFC1.updateFCLayer(fc1) // FC1 - Adam
|
||||
adamFC2.updateFCLayer(fc2) // FC2 - Adam
|
||||
|
||||
// 池化层仍使用 SGD
|
||||
pool1.update(Float32(1.0))
|
||||
pool2.update(Float32(1.0))
|
||||
}
|
||||
}
|
||||
|
||||
// 最终更新 - 使用 Adam
|
||||
adamConv1.updateConvLayer(conv1)
|
||||
adamConv2.updateConvLayer(conv2)
|
||||
adamFC1.updateFCLayer(fc1)
|
||||
adamFC2.updateFCLayer(fc2)
|
||||
pool1.update(Float32(1.0))
|
||||
pool2.update(Float32(1.0))
|
||||
|
||||
let avgLoss = epochLoss / Float32(nTrainSamples)
|
||||
let accuracy = Float32(epochCorrect) / Float32(nTrainSamples) * Float32(100.0)
|
||||
|
||||
if (accuracy > bestAccuracy) {
|
||||
bestAccuracy = accuracy
|
||||
}
|
||||
|
||||
// 输出训练信息
|
||||
if (epoch % 1 == 0 || epoch == nEpochs - 1) {
|
||||
println("Epoch ${epoch + 1}/${nEpochs} - Loss: ${avgLoss}, Accuracy: ${accuracy}%, Adam Step: ${adamConv1.getStep()}")
|
||||
}
|
||||
}
|
||||
|
||||
println("")
|
||||
println("训练完成! 最佳训练准确率: ${bestAccuracy}%")
|
||||
println("总更新步数: ${adamConv1.getStep()}")
|
||||
println("")
|
||||
|
||||
// 测试
|
||||
println("开始测试...")
|
||||
println("")
|
||||
|
||||
var testCorrect = 0
|
||||
var confusion = Array(nClasses, { _ => Array(nClasses, { _ => 0 }) })
|
||||
var testLoss = Float32(0.0)
|
||||
|
||||
for (i in 0..nTestSamples) {
|
||||
let x = testInputs[i]
|
||||
let label = testLabels[i]
|
||||
|
||||
// 前向传播
|
||||
let c1Out = conv1.forward(x, (1, imgSize, imgSize, 1))
|
||||
let p1Out = pool1.forward(c1Out, (1, conv1OutH, conv1OutW, 8))
|
||||
let c2Out = conv2.forward(p1Out, (1, pool1OutH, pool1OutW, 8))
|
||||
let p2Out = pool2.forward(c2Out, (1, conv2OutH, conv2OutW, 16))
|
||||
|
||||
let flatOut = flatten.forward(p2Out, (1, pool2OutH, pool2OutW, 16))
|
||||
|
||||
var fc1Input = Matrix1(1, flatSize)
|
||||
for (j in 0..flatSize) {
|
||||
fc1Input.set(0, j, flatOut.get(j, 0).getOrThrow())
|
||||
}
|
||||
|
||||
let fc1Out = fc1.forward(fc1Input)
|
||||
let fc2Out = fc2.forward(fc1Out)
|
||||
|
||||
let probs = computeSoftmax(fc2Out)
|
||||
|
||||
// 计算损失
|
||||
let probVal = probs.get(0, label).getOrThrow().getValue()
|
||||
testLoss = testLoss + Float32(-log(max(probVal, Float32(1e-10))))
|
||||
|
||||
// 预测
|
||||
let predClass = predictClass(probs)
|
||||
|
||||
confusion[label][predClass] += 1
|
||||
if (predClass == label) {
|
||||
testCorrect += 1
|
||||
}
|
||||
|
||||
// 显示前几个错误分类的样本
|
||||
if (predClass != label && testCorrect < 10) {
|
||||
println("错误分类样本 ${i + 1}:")
|
||||
showSample(x, imgSize, label)
|
||||
let shapeNames = ["圆形", "正方形", "三角形"]
|
||||
println(" 真实标签: ${shapeNames[label]}, 预测标签: ${shapeNames[predClass]}")
|
||||
println("")
|
||||
}
|
||||
}
|
||||
|
||||
let testAcc = Float32(testCorrect) / Float32(nTestSamples) * Float32(100.0)
|
||||
let avgTestLoss = testLoss / Float32(nTestSamples)
|
||||
|
||||
println("测试结果:")
|
||||
println(" 准确率: ${testAcc}% (${testCorrect}/${nTestSamples})")
|
||||
println(" 平均损失: ${avgTestLoss}")
|
||||
println("")
|
||||
|
||||
let classNames = ["圆形 ", "正方形", "三角形"]
|
||||
println("各类别准确率:")
|
||||
for (c in 0..nClasses) {
|
||||
var total = 0
|
||||
for (p in 0..nClasses) {
|
||||
total += confusion[c][p]
|
||||
}
|
||||
if (total > 0) {
|
||||
let acc = Float32(confusion[c][c]) / Float32(total) * Float32(100.0)
|
||||
println(" ${classNames[c]}: ${confusion[c][c]}/${total} (${acc}%)")
|
||||
}
|
||||
}
|
||||
println("")
|
||||
|
||||
println("混淆矩阵:")
|
||||
println(" 预测圆形 预测正方形 预测三角形")
|
||||
for (t in 0..nClasses) {
|
||||
var line = "真实${classNames[t]}:"
|
||||
for (p in 0..nClasses) {
|
||||
line = line + " ${confusion[t][p]} "
|
||||
}
|
||||
println(line)
|
||||
}
|
||||
println("")
|
||||
|
||||
printLine()
|
||||
println("几何图形分类测试完成!")
|
||||
printLine()
|
||||
}
|
||||
+1
-1
@@ -73,7 +73,7 @@ gateway 侧仅做**存在性不变量**检查(显式 `Target-Url` 请求必须
|
||||
|
||||
### `POST /register` / `POST /register/send_code` / `POST /reset/send_code` / `POST /reset_password`
|
||||
|
||||
与 gateway v0.8.0 契约完全一致(详见 `docs/reports/auth_service_plan.md` §5 与黑盒 tests/blackbox/test_gateway.py)。
|
||||
与 gateway v0.8.0 契约完全一致(详见黑盒 tests/blackbox/test_gateway.py)。
|
||||
|
||||
### `POST /logout` / `POST /refresh_token` / `POST /whoami` / `POST /admin/verify`
|
||||
|
||||
|
||||
@@ -56,7 +56,6 @@ captcha2 :8009 (RpcServer)
|
||||
## 状态
|
||||
|
||||
- 生产运行中(services.yaml 拉起),Python captcha 已废弃
|
||||
- 开发文档见 [docs/reports/captcha_cpp_rewrite_dev_doc.md](../docs/reports/captcha_cpp_rewrite_dev_doc.md)
|
||||
- 安全审计见 [docs/security/captcha2_SECURITY_AUDIT_20260809.md](../docs/security/captcha2_SECURITY_AUDIT_20260809.md)
|
||||
|
||||
## 许可证
|
||||
|
||||
+20
-19
@@ -223,22 +223,23 @@ services:
|
||||
depends_on: []
|
||||
sleep_before: 0
|
||||
|
||||
# --- 端口重定向(最后启动) ---
|
||||
- name: redirect80
|
||||
cmd: bash /opt/micro_servers/redirect80/_run.sh
|
||||
args:
|
||||
- "--instance"
|
||||
- "redirect80"
|
||||
- "--log-level"
|
||||
- "info"
|
||||
# --- 配置全部经 CLI 传递(原 confs/configs/redirect80.ini 迁入) ---
|
||||
- "--host"
|
||||
- "0.0.0.0"
|
||||
- "--port"
|
||||
- "80"
|
||||
- "--redirect-port"
|
||||
- "443"
|
||||
- "--redirect-code"
|
||||
- "301"
|
||||
depends_on: [waf]
|
||||
sleep_before: 0
|
||||
# --- 端口重定向(最后启动)---
|
||||
# redirect80 与系统 casaos-gateway 抢占 80 端口导致绑定失败,暂时停用
|
||||
# - name: redirect80
|
||||
# cmd: bash /opt/micro_servers/redirect80/_run.sh
|
||||
# args:
|
||||
# - "--instance"
|
||||
# - "redirect80"
|
||||
# - "--log-level"
|
||||
# - "info"
|
||||
# # --- 配置全部经 CLI 传递(原 confs/configs/redirect80.ini 迁入) ---
|
||||
# - "--host"
|
||||
# - "0.0.0.0"
|
||||
# - "--port"
|
||||
# - "80"
|
||||
# - "--redirect-port"
|
||||
# - "443"
|
||||
# - "--redirect-code"
|
||||
# - "301"
|
||||
# depends_on: [waf]
|
||||
# sleep_before: 0
|
||||
@@ -7,7 +7,8 @@ CREATE TABLE IF NOT EXISTS Users (
|
||||
password_hash TEXT NOT NULL,
|
||||
create_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
last_login DATETIME,
|
||||
is_active BOOLEAN DEFAULT 1
|
||||
is_active BOOLEAN DEFAULT 1,
|
||||
is_admin BOOLEAN NOT NULL DEFAULT 0
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS Notes (
|
||||
@@ -18,14 +19,19 @@ CREATE TABLE IF NOT EXISTS Notes (
|
||||
content TEXT NOT NULL,
|
||||
views INTEGER NOT NULL DEFAULT 0,
|
||||
created_at DATETIME,
|
||||
updated_at DATETIME
|
||||
updated_at DATETIME,
|
||||
is_public BOOLEAN NOT NULL DEFAULT 1,
|
||||
everyone_editable BOOLEAN NOT NULL DEFAULT 1
|
||||
);
|
||||
|
||||
-- 旧版增量
|
||||
-- 旧版增量(存量库请按编号顺序执行对应 patch 脚本)
|
||||
-- ALTER TABLE Notes ADD COLUMN owner VARCHAR(128) DEFAULT 'Unknow';
|
||||
-- ALTER TABLE Notes ADD COLUMN views INTEGER NOT NULL DEFAULT 0; -- patch_001 观看人数列
|
||||
-- ALTER TABLE Notes ADD COLUMN created_at DATETIME; -- patch_002 创建时间
|
||||
-- ALTER TABLE Notes ADD COLUMN updated_at DATETIME; -- patch_002 更新时间
|
||||
-- ALTER TABLE Notes ADD COLUMN views INTEGER NOT NULL DEFAULT 0; -- patch_001 观看人数列
|
||||
-- ALTER TABLE Notes ADD COLUMN created_at DATETIME; -- patch_002 创建时间
|
||||
-- ALTER TABLE Notes ADD COLUMN updated_at DATETIME; -- patch_002 更新时间
|
||||
-- ALTER TABLE Users ADD COLUMN is_admin BOOLEAN NOT NULL DEFAULT 0; -- patch_003 管理员标记
|
||||
-- ALTER TABLE Notes ADD COLUMN is_public BOOLEAN NOT NULL DEFAULT 1; -- patch_003 是否公开(默认公开)
|
||||
-- ALTER TABLE Notes ADD COLUMN everyone_editable BOOLEAN NOT NULL DEFAULT 1; -- patch_003 是否人人可编辑(默认人人可编辑)
|
||||
|
||||
-- Notes 索引:get_notes 按 pid 查子节点、所有权校验按 owner(patch_001 并入)
|
||||
CREATE INDEX IF NOT EXISTS idx_notes_pid ON Notes(pid);
|
||||
|
||||
@@ -0,0 +1,17 @@
|
||||
-- 补丁 003:Users 管理员标记 + Notes 可见性/可编辑性
|
||||
-- 数据库类型: SQLite
|
||||
-- 依赖: create_sqlite3.sql(或等价 schema)已执行,且 patch_001 / patch_002 已按序执行
|
||||
-- 幂等性: ADD COLUMN 为一次性迁移,重复执行会报 "duplicate column name",
|
||||
-- 请按迁移编号顺序执行一次。
|
||||
--
|
||||
-- 语义(默认值向后兼容现有行为):
|
||||
-- Users.is_admin 是否管理员,默认 0(普通用户)
|
||||
-- Notes.is_public 节点是否公开可读,默认 1(公开,与现状一致)
|
||||
-- Notes.everyone_editable 是否任何人都能编辑,默认 1(人人可编辑)
|
||||
|
||||
-- ── Users 新增管理员标记 ─────────────────────────────────────────────
|
||||
ALTER TABLE Users ADD COLUMN is_admin BOOLEAN NOT NULL DEFAULT 0;
|
||||
|
||||
-- ── Notes 新增可见性 / 可编辑性 ──────────────────────────────────────
|
||||
ALTER TABLE Notes ADD COLUMN is_public BOOLEAN NOT NULL DEFAULT 1;
|
||||
ALTER TABLE Notes ADD COLUMN everyone_editable BOOLEAN NOT NULL DEFAULT 1;
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
> 权威架构参考。本文以 `core/services.yaml`(实际运行清单)与 `git HEAD` 为事实来源。
|
||||
> 与根目录 `README.md` 的偏差见文末「§10 与 README 的偏差」。
|
||||
> 最后核对于:2026-08-17。
|
||||
> 最后核对于:2026-08-20。
|
||||
|
||||
---
|
||||
|
||||
@@ -85,7 +85,7 @@ Blackbird 是一套 **进程托管式轻量微服务** 系统(自研 UniServe
|
||||
└─────────────────────────────────────┘
|
||||
|
||||
独立 / 外设(不在主链路):
|
||||
aibox :8765 (Python WebSocket, 独立产品; --once CLI 单次对话 + models 多模型档案含 fallback 失败转移)
|
||||
aibox :8765 (Python WebSocket, 独立产品; --once CLI 单次对话 + models 多模型档案含 fallback 失败转移 + model_config.py 交互式接入点管理)
|
||||
BotOps (aibox 的 TUI 前端)
|
||||
bot4cj (Cangjie + 树莓派 I2C, 硬件)
|
||||
pybot (Python, bot4cj 的迁移版: .mt 解析/录制/实时闭环)
|
||||
@@ -252,6 +252,7 @@ Blackbird 是一套 **进程托管式轻量微服务** 系统(自研 UniServe
|
||||
6. **工作记忆偏差**:✅ 已同步。`ca:8443 → gateway:8445`、`dnfs2`、confs 内嵌 core 均已写入各文档。
|
||||
7. **dnfs2 README 图示滞后**:✅ 已修正(2026-08-10)。`dnfs2/README.md` 架构图与验签描述已更新为 ca:8443(CLI 参数 `--gateway-url` 指向 ca)。
|
||||
8. **memdb / pybot / mtc 未进根 README**:✅ 已回填(2026-08-15)。`memdb:8011`、`pybot`(bot4cj Python 迁移版)、`mtc`(Python DSL → .mt 编译器,2026-08-15 落地)连同 `actpred` 已入根 `README.md` 模块表与端口一览;本文件与 MODULES.md 已同步。
|
||||
9. **aibox 第五轮审计(Batch K)与 model_config.py**:✅ 已同步(2026-08-20)。根 README aibox 小节审计轮次更新为五轮(E/F/G/H/I/K,单测 490 项全过)并增补 `model_config.py` 交互式接入点管理;MODULES.md 同步 aibox 行与文档计数;aibox/README.md 更新日志补 v3.6 条目(J1~J16 + Batch K 波次 1~5 落地)。
|
||||
|
||||
---
|
||||
|
||||
|
||||
+2
-2
@@ -44,7 +44,7 @@
|
||||
|
||||
| 模块 | 语言 | 端口 | 职责 |
|
||||
|------|------|------|------|
|
||||
| aibox | Python | 8765 | 独立 AI 智能体系统(DeepSeek + WebSocket),自带技能系统/记忆/任务调度/前端与语音客户端;`--once` CLI 单次对话(stdout/stderr 分流),`models` 配置节支持多 OpenAI 兼容端点按档案名切换 |
|
||||
| aibox | Python | 8765 | 独立 AI 智能体系统(DeepSeek + WebSocket),自带技能系统/记忆/任务调度/前端与语音客户端;`--once` CLI 单次对话(stdout/stderr 分流),`models` 配置节支持多 OpenAI 兼容端点按档案名切换,`model_config.py` 交互式接入点管理(列表/添加/删除/连通性验证,Key 仅记环境变量名不落盘) |
|
||||
| BotOps | Python/Textual | — | aibox 的终端 TUI 前端,WebSocket 连接 aibox,Markdown 增量流式渲染、思考模式、任务侧栏、自动重连 |
|
||||
| actpred | Python/PyTorch | — | 动作预测(图像 + 动作历史 → 预测下一个动作),bot4cj 的「大脑」配套;支持从 `.mt` 脚本预训练 |
|
||||
|
||||
@@ -77,7 +77,7 @@
|
||||
| 模块 | 说明 |
|
||||
|------|------|
|
||||
| tests | 黑盒测试套件(`test_db.py`/`test_gateway.py`/`test_waf.py`/`test_dnfs.py`/`test_captcha.py`/`test_smtp.py`/`test_ai_review.py`/`test_static.py`/`test_core_confs.py`),纯标准库 HTTP/TCP 用户旅程模拟,经 `test_runner.py` 统一调度(详见 [tests/blackbox/README.md](../tests/blackbox/README.md)) |
|
||||
| docs | 文档聚合(security/ 审计、reviews/ 审查、reports/ 分析),76 篇 md |
|
||||
| docs | 文档聚合(security/ 审计、reviews/ 审查、reports/ 分析),64 篇 md + 1 篇 html |
|
||||
|
||||
## 九、根级构建与部署脚本
|
||||
|
||||
|
||||
@@ -1,101 +0,0 @@
|
||||
# actpred 动作归一化与逐列权重开发方案
|
||||
|
||||
> 文档类型:开发方案
|
||||
> 模块:actpred(联动 pybot 数据契约)
|
||||
> 作者:自动生成
|
||||
> 日期:2026-08-14
|
||||
> 前置实验:`pybot_actpred_dataset_training_eval_20260814.md`
|
||||
|
||||
---
|
||||
|
||||
## 1. 问题陈述
|
||||
|
||||
闭环评估(R4)暴露:z-score 归一化把 18 列方差拉平,隐式实现"全列等权",导致**舵机精度 ↔ spk.F 频率跟踪的权衡**——未归一化模型舵机准但 spk.F 崩(21k MSE),归一化模型 spk.F 改善 3× 但简单动作恶化 1.3-25×。
|
||||
|
||||
本质:**缩放方式(数值稳定性)与列重要性(优化方向)是两个正交维度**,当前被隐式耦合。方案目标:固定物理范围归一化(数值稳定、免 stats 文件)+ 显式逐列权重(舵机优先),两者解耦。
|
||||
|
||||
## 2. 方案设计
|
||||
|
||||
### 2.1 固定物理范围 min-max 归一化(actpred 侧)
|
||||
|
||||
位置决策:**归一化保留在 actpred 侧,pybot 继续记录原始量纲**。理由:
|
||||
|
||||
- **raw at rest**:episode 是数据资产,原始角度/频率可复用于逆运动学、碰撞检测等分析,归一化不可逆
|
||||
- **反归一化链条**:若 pybot 侧归一化,`--loop` 执行前需反归一化,min/max 表要同时存在于 recorder 与 loop 两端;现状 `predict()` 返回原始量纲,pybot 无感知,边界更干净
|
||||
- **代码库先例**:IMU 即"pybot 记录原始、actpred `imu_norm.py` 按物理量程归一化",动作走同一条路保持一致
|
||||
- **pybot 零依赖承诺**:归一化逻辑入 recorder 会加深 pybot 对 numpy 的依赖面
|
||||
|
||||
固定范围常量表(新增 `actpred/action_norm.py`,类比 `imu_norm.py`):
|
||||
|
||||
| 列 | 物理范围 | 依据 |
|
||||
|---|---|---|
|
||||
| 0-9 舵机角 | **[0, 180]** | `pybot/hd/servo.py` `MAX_ANGLE=180.0`(校准后全行程 0-180°) |
|
||||
| 10 spk.F | [0, 5000] | `@Speaker max_f=5000` |
|
||||
| 11 spk.delay | [0, 2000] | 实用上限(脚本用 10-500ms) |
|
||||
| 12-17 保留位 | 恒 0 | 不参与归一化,损失权重 0 |
|
||||
|
||||
min-max 映射:`x_norm = (x - lo) / (hi - lo)` → [0,1](或映射到 [-1,1],与 IMU 归一化输出区间一致,待实现时定)。
|
||||
|
||||
优点:免 action_stats.npy(消除 stats 路径类 bug 的整类风险)、零状态部署、数值区间有界。风险:min-max 对离群值敏感——若真实数据某列越界,需 clamp 到 [0,1] 并告警(防御口径与 `_action_bounds` 一致)。
|
||||
|
||||
可选细化(后续):按各舵机 `|max|`(s1=90°、s2=180°…)分列归一化,90° 舵机能用满 [0,1],分辨率更高;需要读取 @Servo 块的 max 列构建表。
|
||||
|
||||
### 2.2 逐列加权损失
|
||||
|
||||
```python
|
||||
COLUMN_WEIGHTS = np.array([
|
||||
1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, # 0-9 舵机(安全关键)
|
||||
0.2, 0.2, # 10-11 spk.F / spk.delay(声音与节奏,不影响运动)
|
||||
0.0, 0.0, 0.0, 0.0, 0.0, 0.0, # 12-17 保留位(恒 0,纯噪声)
|
||||
], dtype=np.float32)
|
||||
# loss = mean(w * (pred - target)^2) / mean(w)
|
||||
```
|
||||
|
||||
实施点:
|
||||
|
||||
- `train.py` / `pretrain_mt.py` 的 `train_one_epoch` / `validate`:MSE 处乘权重向量(权重存 DataConfig,随 checkpoint 序列化,推理端无需感知)
|
||||
- 权重可经 CLI `--column-weights` 覆盖(默认用上表)
|
||||
- 训练集/验证集用同一权重;`evaluate.py` 的逐列误差报告不受影响(报告仍按无权重口径,便于对照)
|
||||
|
||||
### 2.3 归一化与权重的组合矩阵
|
||||
|
||||
| 配置 | 数值稳定性 | 舵机精度 | spk.F 跟踪 | 预期适用 |
|
||||
|---|---|---|---|---|
|
||||
| raw + 等权(模型 A) | 差(loss 数百量级) | 好 | 崩 | 仅演示/对照 |
|
||||
| z-score + 等权(模型 B) | 好 | 差 | 好 | 现状 |
|
||||
| **固定范围 + 加权(本方案)** | 好 | 好 | 中(w=0.2 保留梯度) | **目标** |
|
||||
|
||||
## 3. 配套改进(同批实施)
|
||||
|
||||
### 3.1 R4 评估器修正(指标失效修复)
|
||||
|
||||
1. **oob_rate 恒值列退化**:span<阈值(如 1e-3)的列不参与越界判定;或改用固定物理界(舵机 [0,180]、spk.F [0,5000])。推荐后者——与 §2.1 常量表共用,且"越物理界"比"越经验界"更符合安全语义
|
||||
2. **error_growth 长度混叠**:按 episode 归一化后聚合(每 episode 自身首末段比值取中位数),避免短 episode 尖峰主导曲线
|
||||
|
||||
### 3.2 R1-B rollout + scheduled sampling
|
||||
|
||||
闭环误差累积源于 teacher forcing 训练与自预测推理的不匹配。已实现的 `--rollout-len K --ss-prob p` 直接对标:建议 `K=4, p=0.5`,与加权损失同批训练对照。
|
||||
|
||||
### 3.3 数据扩展
|
||||
|
||||
- **实机录制**(根本瓶颈):dry-run 帧为噪声、IMU 全零,视觉/姿态通路无真实信号。实机 `--record` 后复跑本实验
|
||||
- 验证集 1 个 episode 选模噪声大:val_split 0.15→0.3
|
||||
- 补充站立/保持类长序列(当前 std 类仅 1-81 步)
|
||||
|
||||
## 4. 实施步骤
|
||||
|
||||
| 步骤 | 内容 | 产出 |
|
||||
|---|---|---|
|
||||
| 1 | 新增 `actpred/action_norm.py`:固定范围常量表 + normalize/denormalize + clamp 防御 | 独立纯 numpy 模块 + 单测 |
|
||||
| 2 | DataConfig 增 `normalize_mode`(`zscore`/`fixed_range`)与 `column_weights`;dataset/train 接入 | 训练链路支持 + 单测 |
|
||||
| 3 | `evaluate.py` oob 改物理界、error_growth 按 episode 归一化 | 评估器修正 + 单测 |
|
||||
| 4 | 训练对照:固定范围+加权 vs z-score+等权(同种子、同数据) | 模型 C/D checkpoint |
|
||||
| 5 | 复跑 R4 闭环评估,验证权衡解耦 | 对比报告 |
|
||||
|
||||
步骤 1-3 可并行;步骤 4-5 依赖前置完成。预计步骤 1-3 改动集中在 `actpred/`(`action_norm.py` 新增、`config.py`/`dataset.py`/`train.py`/`evaluate.py` 小改),不影响 pybot 与部署链(归一化在推理端内部完成,`predict()` 输出仍为原始量纲)。
|
||||
|
||||
## 5. 决策记录
|
||||
|
||||
- **归一化位置**:讨论过"移入 pybot 按列 min/max 归一化"。结论:不采纳(raw at rest、反归一化链条变长、IMU 先例不一致),但采纳其"固定物理范围"内核,落在 actpred 侧常量表实现
|
||||
- **舵机范围**:±180° 为误述,实际全行程 0-180°(`servo.py MAX_ANGLE=180.0`),归一化表按 [0,180]
|
||||
- **spk.F 对数尺度**:备选未定。若加权后 spk.F 跟踪仍不足,再考虑 log 变换(频率感知是对数的,10→1000 跨两个数量级)
|
||||
@@ -4,8 +4,7 @@
|
||||
> 模块:actpred(hailo_infer.py cv2.LUT 快路径)
|
||||
> 作者:自动生成
|
||||
> 日期:2026-08-17
|
||||
> 前置:`actpred_hailo_preprocess_optimization_plan_20260817.md`(优化方案,已实施)、
|
||||
> `actpred_inference_speed_benchmark_20260813.md`(旧基准)
|
||||
> 前置:`actpred_inference_speed_benchmark_20260813.md`(旧基准)
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -5,7 +5,6 @@
|
||||
> 作者:自动生成
|
||||
> 日期:2026-08-14
|
||||
> 前置实验:`pybot_actpred_dataset_training_eval_20260814.md`(模型 A/B)
|
||||
> 开发方案:`actpred_action_norm_weighting_plan_20260814.md`(本实验为其步骤 4)
|
||||
> 后续实验:模型 D(fixed_range + 加权 + R1-B rollout),训练中
|
||||
|
||||
---
|
||||
|
||||
@@ -1,154 +0,0 @@
|
||||
# actpred Hailo 推理预处理优化:cv2.LUT 查表合并归一化与量化(已实施)
|
||||
|
||||
> 文档类型:开发方案(已实施)
|
||||
> 模块:actpred(部署链 hailo_infer.py)
|
||||
> 作者:自动生成
|
||||
> 日期:2026-08-17(原型实测 + 集成实施)
|
||||
> 前置:`actpred_inference_speed_benchmark_20260813.md`(速度基准)、
|
||||
> 2026-08-17 新 640×480 数据集重测(本方案 §2 数据)
|
||||
|
||||
> 📌 **版本历程**
|
||||
> 1. 初版:numpy 查表法为核心,预期端到端 18-20ms
|
||||
> 2. 原型实测(§3.4):**numpy 高级索引查表在 Pi5 上无加速**(fancy
|
||||
> indexing 的 gather 开销与逐元素 float 运算相当),预期不成立;
|
||||
> 合并系数乘加仅 -5ms;**cv2.LUT 实测 15.8× 加速**(§3.4)
|
||||
> 3. **已实施(§5)**:`hailo_infer.py` 集成 cv2.LUT 快路径(可选依赖,
|
||||
> 无 cv2 自动回退 float 路径),端到端 52.4 → **16.0 ms**,输出与
|
||||
> float 路径逐点一致(max_diff = 0)
|
||||
|
||||
---
|
||||
|
||||
## 1. 背景与动机
|
||||
|
||||
2026-08-17 用新录制的 640×480 数据集(免 resize 训练样本)重测 HEF 推理速度:
|
||||
|
||||
| 阶段 | 640×480 输入 | 224×224 输入(旧数据) |
|
||||
|---|---|---|
|
||||
| 端到端 predict | **52.4 ms** | 60.0 ms |
|
||||
| 预处理 | **35.9 ms** | 43.1 ms |
|
||||
| NPU run | 12.5 ms | 12.9 ms |
|
||||
| 后处理 | 0.6 ms | 0.6 ms |
|
||||
|
||||
输入尺寸已与模型一致(resize 变 no-op),**预处理仍是端到端延迟的主导项**(68%)。
|
||||
旧基准报告(2026-08-13)已记录同类问题:1080p 采集的预处理缩放约 45ms 主导端到端延迟,
|
||||
当时通过把默认分辨率改回 640×480 缓解;本次分析证明 640×480 下 Python 端预处理
|
||||
仍有约 36ms 的优化空间——但空间比初版估计小(见 §3.4 实测)。
|
||||
|
||||
## 2. 瓶颈分析(100 次均值分段计时)
|
||||
|
||||
| 子步骤 | 耗时 | 说明 |
|
||||
|---|---|---|
|
||||
| `_frame_float`(3 帧) | **27.2 ms** | PIL fromarray/convert + resize(同尺寸 no-op 仍有开销)+ `(x/255-mean)/std` float 运算(92 万像素×3 通道) |
|
||||
| `create_bindings` + 帧量化绑定 | **9.1 ms** | `_bind_frames`:float 归一化帧 → `round(x/q.scale+q.zp)` 再量化回 uint8 |
|
||||
| history/pred 绑定 | 0.07 ms | 可忽略 |
|
||||
|
||||
**关键冗余**:当前流程先做 float 归一化(`_frame_float`),随后又量化回 uint8(`_bind_frames`)——
|
||||
中间 float 帧是纯中间态,归一化与量化可以**合并为一步**。
|
||||
|
||||
## 3. 优化方案
|
||||
|
||||
### 3.1 主方案:合并系数乘加(纯 numpy,无新依赖)
|
||||
|
||||
HEF 帧输入是 uint8 量化值:`q = round(((x/255 - mean)/std) / q_scale + q_zp)`。
|
||||
把 4 步运算合并为**一次乘加**(系数预计算,按通道):
|
||||
|
||||
```
|
||||
q = round(x * A[c] + B[c])
|
||||
A[c] = 1 / (255 * std[c] * q_scale)
|
||||
B[c] = -mean[c] / (std[c] * q_scale) + q_zp
|
||||
```
|
||||
|
||||
- 现流程 4 次全数组运算(`/255`、`-mean`、`/std`、`/q_scale+q_zp`)→ 1 次乘加
|
||||
- 数值完全等价(实测 max_diff = 0 LSB,见 §3.4)
|
||||
- 直通模式(`_normalize_frames=False`)同理:`A=1/q_scale, B=q_zp`
|
||||
|
||||
### 3.2 附带优化(与主方案正交)
|
||||
|
||||
1. **免 PIL 路径**:输入已是 np.ndarray 时直接运算,跳过 `Image.fromarray/convert`
|
||||
2. **resize no-op 跳过**:`image.size == (W,H)` 时不调用 `resize()`
|
||||
3. **缓冲复用**:`np.multiply/np.add` 带 `out=` 参数,减少临时数组分配
|
||||
|
||||
### 3.3 已实施快路径:cv2.LUT(可选依赖,无 cv2 自动回退)
|
||||
|
||||
`cv2.LUT(img, lut)` 是 C 实现的 uint8→uint8 查表,单次遍历,无 float 中间态——
|
||||
把"归一化+量化"合并为一步。集成要点(`hailo_infer.py`):
|
||||
|
||||
- 惰性导入 cv2(`HAS_CV2`);**无 cv2 或非逐帧输入布局(ncfhw/nhwc_fused)自动回退 float 路径**
|
||||
- 加载时按每帧 `quant_info` 构建 256 项表(归一化/直通两种,`_normalize_frames` 判定)
|
||||
- **OpenCV 5.0 的 cv2.LUT 只接受单通道表**((256,3) 会断言失败),逐通道应用后 `cv2.merge`
|
||||
- 输入尺寸与模型不一致时先 PIL resize(仅此情况走 PIL)
|
||||
- 依赖安装:`opencv-python-headless`(5.0.0,清华镜像;官方源过慢)
|
||||
- 注:本方案实测后发现**合并系数乘加(§3.1)未实施**——cv2.LUT 收益已远超,
|
||||
避免冗余代码;§3.1 保留作为无 cv2 环境下的潜在备选
|
||||
|
||||
### 3.4 原型实测评估(2026-08-17,100-200 次均值,640×480×3 帧)
|
||||
|
||||
| 方案 | 3 帧耗时 | 对比现 float 28.5ms | 数值差异 |
|
||||
|---|---|---|---|
|
||||
| 现 float 路径 | 28.5 ms | 基线 | — |
|
||||
| numpy 高级索引查表 | 26.6 ms | **无加速(1.0x)** | 0 LSB |
|
||||
| 合并系数乘加 | 23.2 ms | 快 19%(-5ms) | 0 LSB |
|
||||
| **cv2.LUT(逐通道)** | **1.68 ms** | **快 15.8×** | **0 LSB** |
|
||||
|
||||
**结论**:
|
||||
- numpy fancy indexing 的 gather 开销与逐元素 float 运算相当,查表法**无收益**
|
||||
- 纯 numpy 可达的最佳是合并乘加(-5ms),受限于 92 万像素 float32 数组的内存吞吐
|
||||
- cv2.LUT 的 C 查表把帧预处理降一个数量级,且数值完全等价(0 LSB)
|
||||
|
||||
### 3.5 实测收益(已实施,集成后 NPU 实测)
|
||||
|
||||
| 指标 | 优化前 | 优化后(cv2.LUT) |
|
||||
|---|---|---|
|
||||
| 帧预处理(3 帧) | 26.5 ms | **3.0 ms(分段实测含绑定)** |
|
||||
| 端到端 predict | 52.4 ms | **16.1 ms(3.3×)**(200 次均值,见测试报告) |
|
||||
| 输出一致性 | — | 与 float 路径 max_diff = **0.0000**(7 类输入 + 8 seed 随机) |
|
||||
| HEF vs torch | 0.50° | 0.30°(4 角复验 PASS) |
|
||||
| vs PyTorch CPU | 2.1× | **6.7×**(16.1 vs 107.9 ms) |
|
||||
|
||||
闭环 10Hz(100ms 预算)占用率 52% → **16%**,20Hz 32%,30Hz(33ms)48%。
|
||||
|
||||
> 完整验收数据见 `actpred_cv2_lut_benchmark_report_20260817.md`(测试报告)。
|
||||
|
||||
## 4. 一致性约束(关键)
|
||||
|
||||
- **训练侧不动**:`dataset.py` 的 `Resize+ToTensor+Normalize` 保持 float 路径——
|
||||
合并乘加只是推理侧对"归一化+量化"的合并实现,数值等价(实测 0 LSB)
|
||||
- **三后端口径**:`predict.py`(PyTorch,float 输入不量化)、`onnx_infer.py`(ORT,float)
|
||||
不受影响;优化仅作用于 `hailo_infer.py` 的 uint8 量化输入路径
|
||||
- **帧输入判定不变**:`_normalize_frames`(量化网格负零偏判定)决定走
|
||||
归一化还是直通,逻辑与现状一致
|
||||
- 旧 5 输入 / 旧版单图像 HEF 布局同样适用(按各自 quant_info 算系数)
|
||||
|
||||
## 5. 实施步骤(已完成 ✅)
|
||||
|
||||
| 步骤 | 内容 | 状态 |
|
||||
|---|---|---|
|
||||
| 1 | `hailo_infer.py`:加载时按每帧 quant_info 构建 256 项 LUT(归一化/直通两种) | ✅ `_build_frame_luts` |
|
||||
| 2 | `predict` 的 nhwc_frames 分支走 cv2.LUT 快路径(尺寸不匹配时 PIL resize);无 cv2 回退 float | ✅ `_frame_lut_quant` + `HAS_CV2` |
|
||||
| 3 | NPU 实测:640×480 数据集 100 次均值 | ✅ 16.0 ms(52.4 → 16.0) |
|
||||
| 4 | `_npu_test.py` 4 角 + HEF vs torch 对比复验 | ✅ PASS,0.30°(≤1° 预算) |
|
||||
| 5 | 全套单测(TestHailoPredictor 等) | ✅ 188 tests OK |
|
||||
| 6 | 依赖安装:opencv-python-headless 5.0.0(清华镜像) | ✅ venv 已装 |
|
||||
|
||||
## 6. 风险与对策
|
||||
|
||||
| 风险 | 对策 |
|
||||
|---|---|
|
||||
| 查表取整与 float 路径存在差异 | 实测 0 LSB(表项 float32 与现路径同精度);NPU 输出 vs float 路径 max_diff=0,vs torch 0.30° |
|
||||
| 无 cv2 环境 | 惰性导入 + 自动回退 float 路径,行为与优化前完全一致 |
|
||||
| OpenCV 5.0 多通道 LUT 断言失败 | 逐通道单表 + `cv2.merge`(实测无性能损失) |
|
||||
| 非 nhwc_frames 布局(旧 HEF) | 仅该布局启用 cv2 路径,其余走原 float 路径 |
|
||||
| 未来输入尺寸变更 | LUT 与尺寸无关(逐像素映射);尺寸不一致时才走 PIL resize |
|
||||
|
||||
## 7. 决策记录
|
||||
|
||||
- **查表法(初版核心)实测证伪**:numpy 高级索引 gather 在 Pi5 上与逐元素 float
|
||||
运算同速(26.6 vs 28.5ms),"查表降为内存索引即快"的假设不成立——教训:
|
||||
性能方案须原型实测,不能凭直觉估计
|
||||
- **最终方案选 cv2.LUT**:C 实现查表把帧预处理从 26.5ms 降到 1.7ms(15.8×),
|
||||
端到端 52.4 → 16.0ms;数值 0 LSB 等价。合并系数乘加(-5ms)收益被碾压,
|
||||
不再实施(§3.1 保留为无 cv2 环境的备选)
|
||||
- **依赖取舍**:opencv-python-headless 作为可选依赖(惰性导入 + 自动回退),
|
||||
有 cv2 环境收益 3.3×,无 cv2 环境行为不变——收益远大于依赖成本
|
||||
- **训练侧不改**:训练用 float 归一化(增强/数值精度),推理侧查表实现——两侧
|
||||
口径在"归一化值"层面一致,量化差异由部署链固有预算吸收(HEF vs torch 0.30° 余量充足)
|
||||
@@ -1,82 +0,0 @@
|
||||
# actpred 多头 one-hot 分类输出开发方案
|
||||
|
||||
> 文档类型:开发方案
|
||||
> 模块:actpred(联动 pybot 数据契约)
|
||||
> 作者:自动生成
|
||||
> 日期:2026-08-14
|
||||
> 前置实验:`pybot_actpred_dataset_training_eval_20260814.md`、
|
||||
> `actpred_fixed_range_weighted_experiment_20260814.md`(模型 A-E 系列)
|
||||
|
||||
---
|
||||
|
||||
## 1. 背景与动机
|
||||
|
||||
模型 A-E 系列实验定位到两个结构性难题:
|
||||
|
||||
1. **尺度放大**:fixed_range 线性 min-max 下 spk.F 的 1 个归一化单位 = 5000Hz 原始量,闭环反馈中误差滚雪球(模型 C/D 发散根因之一)
|
||||
2. **精度错配**:常用值(10Hz → 0.002)挤在线性尺度的低端,[0,1] 精度预算被浪费(log_range 正在验证的对数尺度解法)
|
||||
|
||||
多头 one-hot 是本问题的**结构性解法**:把 18 维连续回归换成 18 组独立分类头,每组输出一个 one-hot 分布。分类天然消除尺度放大(每箱等距)与物理界问题(箱只覆盖物理范围),且多列 softmax 是 NPU(Hailo)最擅长的计算形态。
|
||||
|
||||
## 2. 方案设计
|
||||
|
||||
### 2.1 分箱设计(每组 one-hot 的维度)
|
||||
|
||||
| 列组 | 分箱方式 | 箱数/列 | 依据 |
|
||||
|---|---|---|---|
|
||||
| 舵机 0-9(10 列) | 均匀 **1°/箱** | 180 | 匹配数据粒度:微动作插值 0.75°/步,5° 箱会把相邻步量化进同一箱、毁掉插值结构 |
|
||||
| spk.F(10) | 对数分箱 | 32 | 频率对数感知,低端聚集;log10(1+x) 均匀切 |
|
||||
| spk.delay(11) | 对数分箱 | 32 | 同左 |
|
||||
| 保留位 12-17 | 无头 | 0 | 恒 0,不参与 |
|
||||
|
||||
**输出总量**:10×180 + 32 + 32 = **1864 logits**(当前 18)。最后层 128→1864 ≈ 238K 参数,模型 409K → ~650K,仍属轻量。
|
||||
|
||||
### 2.2 各层改动
|
||||
|
||||
| 文件 | 改动 |
|
||||
|---|---|
|
||||
| `actpred/bins.py`(新增) | 分箱表 + `encode(x)→one-hot` / `decode(p)→连续值`(含对数箱);纯 numpy,供推理/部署复用 |
|
||||
| `model.py` | PredictionHead 输出 1864 维 → reshape (B, 18, bins) → 逐组 softmax |
|
||||
| `dataset.py` | target 从 (B,18) 连续值 → 每组 one-hot;**软标签**:真值高斯扩散到 ±1-2 箱 |
|
||||
| `train.py` | 损失 = Σ 列权重 × CE(组输出, one-hot 真值);列权重沿用 default_column_weights |
|
||||
| `predict.py` | 输出 decode:**期望值** `Σ pᵢ·箱中心`(推荐)或 argmax 箱中心 → 18 维原始量纲 |
|
||||
| `evaluate.py` | 无感——predict 返回仍是 18 维原始量纲 |
|
||||
| `export_deploy` / Hailo | 多组 softmax 为 NPU 友好形态 |
|
||||
|
||||
### 2.3 解码策略(执行时如何取连续值)
|
||||
|
||||
| 方式 | 行为 | 评估 |
|
||||
|---|---|---|
|
||||
| **期望值**(推荐) | `Σ pᵢ × 箱中心` | 平滑连续,分辨率超越箱宽(1° 箱可出 46.4°);单峰假设 |
|
||||
| argmax | 取最大概率箱中心 | 确定性,量化到 1° |
|
||||
| 按分布采样 | 随机抽箱 | 探索用,后续离线 RL 可复用 |
|
||||
|
||||
## 3. 关键决策点与理由
|
||||
|
||||
1. **舵机 1°/箱而非 5°/箱**:录制数据的微动作插值粒度为 0.75°,箱宽必须匹配数据粒度;5° 箱会把 45.0°/45.75° 压进同一箱,模型学不到特意录制的插值结构
|
||||
2. **软标签**:804 样本对 1864 logits 偏少,多数箱极少被访问。真值 one-hot 高斯扩散到邻近箱,未访问箱也能获得梯度
|
||||
3. **期望值解码补序数信息**:one-hot 丢失"45° 和 46° 相邻"先验,期望值解码在推理端部分恢复;训练端靠软标签缓解
|
||||
4. **pybot 契约零变化**:predict 输出仍是 18 维原始量纲,闭环与录制数据无感知
|
||||
5. **与现有归一化链路的兼容**:one-hot 方案独立于 zscore/fixed_range/log_range——分箱编码替代归一化,不叠加使用
|
||||
|
||||
## 4. 实施步骤
|
||||
|
||||
| 步骤 | 内容 | 产出 |
|
||||
|---|---|---|
|
||||
| 1 | `bins.py`:分箱表 + encode/decode + 对数箱 + 软标签工具 | 独立纯 numpy 模块 + 单测 |
|
||||
| 2 | `dataset.py`:目标 one-hot 化(软标签可开关) | 数据链路 + 单测 |
|
||||
| 3 | `model.py` + `train.py`:多头 softmax + 加权 CE | 模型与训练 + 单测 |
|
||||
| 4 | `predict.py`:期望值/argmax 解码集成 + 冒烟 | 推理链路 + 单测 |
|
||||
| 5 | 训练模型 F,与 A-E 闭环五方对比 | 对比报告 |
|
||||
|
||||
## 5. 与其他方案的定位
|
||||
|
||||
- **模型 E(log_range)正在训练**:若对数尺度已解决精度错配与闭环发散,one-hot 降级为可选优化(分类头天然有界、免归一化常量表的优势仍在)
|
||||
- 若 E 仍失败:one-hot 是下一个强候选(结构性消除尺度放大)
|
||||
- **互补关系**:one-hot 方案与逐列权重、R1-B rollout 正交,可叠加(rollout 治 exposure bias 的结论不变,任何损失形式都需要它)
|
||||
|
||||
## 6. 决策记录
|
||||
|
||||
- **箱宽**:初拟 5°/箱(保守,顾虑数据量),讨论后定为 1°/箱——数据插值粒度 0.75° 决定箱宽下限;spk 两列对数 32 箱
|
||||
- **单头多选 vs 多头单选**:softmax 单头在数学上是单选择分布;"多组 one-hot" = 18 个独立 softmax 头并行,每列一个选择
|
||||
- **多峰分布**:当前数据单峰为主,期望值解码够用;若未来出现双峰(如跨脚本姿态冲突),升级为逐箱 sigmoid 或 MDN
|
||||
@@ -1,116 +0,0 @@
|
||||
# actpred 训练时安全图像增强
|
||||
|
||||
> 文档类型:特性报告
|
||||
> 模块:actpred
|
||||
> 日期:2026-08-13
|
||||
|
||||
---
|
||||
|
||||
## 1. 背景
|
||||
|
||||
pybot dry-run 录制的 episode 数据样本量有限(11 个脚本 → 12 个 episode → ~728 样本),训练容易过拟合。直接增大数据集需要更多 `.mt` 脚本或更多轮次录制,成本较高。
|
||||
|
||||
图像增强(data augmentation)是低成本扩充样本的标准手段,但机器人动作预测场景有特殊约束:**空间变换(翻转/旋转/裁剪)会改变图像的语义坐标,而动作标签未随之变换,导致训练对错误**。例如水平翻转后,左右舵机角度应互换,但标签不变,形成错误训练信号。
|
||||
|
||||
## 2. 方案
|
||||
|
||||
### 2.1 增强范围
|
||||
|
||||
仅启用**颜色空间变换**(ColorJitter),不涉及任何空间变换:
|
||||
|
||||
| 变换 | 参数 | 说明 |
|
||||
|------|------|------|
|
||||
| 亮度 | `±20%`(factor ∈ [0.8, 1.2]) | 模拟不同光照条件 |
|
||||
| 对比度 | `±20%`(factor ∈ [0.8, 1.2]) | 模拟不同摄像头增益 |
|
||||
| 饱和度 | `±20%`(factor ∈ [0.8, 1.2]) | 模拟不同色温环境 |
|
||||
| 色调 | `±5%`(factor ∈ [-0.05, 0.05]) | 轻微色偏,避免模型过依赖颜色 |
|
||||
|
||||
**不使用的变换**(因需同步变换动作标签):
|
||||
|
||||
| 变换 | 原因 |
|
||||
|------|------|
|
||||
| 水平/垂直翻转 | 左右舵机角度需互换,否则标签错误 |
|
||||
| 旋转 | 重力方向变化,姿态语义改变 |
|
||||
| 裁剪/平移 | 视野变化导致动作-画面对应关系改变 |
|
||||
| 透视变换 | 改变空间关系 |
|
||||
|
||||
### 2.2 同一样本内帧间参数一致性
|
||||
|
||||
一个训练样本包含 `frame_len` 帧连续图像(默认 3 帧,间隔 10ms)。真实场景中连续帧的光照条件不变,因此**同一样本内所有帧共享同一组增强参数**,而非每帧独立随机:
|
||||
|
||||
```python
|
||||
if self.augment:
|
||||
bf = random.uniform(0.8, 1.2) # 亮度因子,所有帧共用
|
||||
cf = random.uniform(0.8, 1.2) # 对比度因子
|
||||
sf = random.uniform(0.8, 1.2) # 饱和度因子
|
||||
hf = random.uniform(-0.05, 0.05) # 色调因子
|
||||
for p in image_paths:
|
||||
image = Image.open(p).convert("RGB")
|
||||
if self.augment:
|
||||
image = TF.adjust_brightness(image, bf)
|
||||
image = TF.adjust_contrast(image, cf)
|
||||
image = TF.adjust_saturation(image, sf)
|
||||
image = TF.adjust_hue(image, hf)
|
||||
frames.append(self.transform(image))
|
||||
```
|
||||
|
||||
不同样本(不同 `__getitem__` 调用)使用独立随机参数,保证整个数据集的增强多样性。
|
||||
|
||||
### 2.3 自适应 frame_len
|
||||
|
||||
增强逻辑在 `for p in image_paths` 循环前生成参数、循环内逐帧施加。`image_paths` 的长度由 `frame_len` 决定(`image_files[t - frame_len + 1 : t + 1]`),循环自动适配。**`frame_len` 为 1、2、3、5、10 等任意值时均无需修改代码**。
|
||||
|
||||
### 2.4 训练/验证隔离
|
||||
|
||||
| 数据集 | 增强 | 说明 |
|
||||
|--------|------|------|
|
||||
| 训练集 | `augment=cfg.train.augment` | CLI `--augment` 控制 |
|
||||
| 验证集 | `augment=False`(默认) | 验证 loss 反映真实分布表现 |
|
||||
|
||||
推理端(`predict.py` / `onnx_infer.py` / `hailo_infer.py`)不涉及增强,使用原始图像。
|
||||
|
||||
## 3. CLI 接口
|
||||
|
||||
```bash
|
||||
# 训练时启用图像增强
|
||||
bash _run.sh train --data-dir ./data --epochs 100 --augment
|
||||
|
||||
# 显式关闭(默认关闭)
|
||||
bash _run.sh train --data-dir ./data --epochs 100 --no-augment
|
||||
```
|
||||
|
||||
`--augment` 使用 `BooleanOptionalAction`,支持 `--augment` / `--no-augment` 双向开关。
|
||||
|
||||
## 4. 验证
|
||||
|
||||
### 4.1 功能验证
|
||||
|
||||
```
|
||||
dataset size: 728 samples
|
||||
augment diff (should be >0): 0.3767 # 增强前后图像有差异
|
||||
two augment calls diff (should be >0): 0.2626 # 两次增强参数不同(随机性)
|
||||
action_history diff (should be 0): 0.0 # 动作标签不受影响
|
||||
target diff (should be 0): 0.0
|
||||
```
|
||||
|
||||
### 4.2 端到端训练
|
||||
|
||||
`--augment --epochs 3`(64×64,batch=8)训练正常完成,loss 正常下降。
|
||||
|
||||
### 4.3 回归测试
|
||||
|
||||
全量单测 90/90 通过,无回归。
|
||||
|
||||
## 5. 改动文件
|
||||
|
||||
| 文件 | 改动 |
|
||||
|------|------|
|
||||
| `actpred/dataset.py` | `ActionDataset` 新增 `augment` 参数;`__getitem__` 在 Resize 前施加 ColorJitter(`torchvision.transforms.functional`) |
|
||||
| `actpred/config.py` | `TrainConfig` 新增 `augment: bool = False`;`build_argparser` 新增 `--augment`;`apply_overrides` 新增映射 |
|
||||
| `actpred/train.py` | 训练集构造传 `augment=cfg.train.augment`,验证集默认不增强 |
|
||||
|
||||
## 6. 后续方向
|
||||
|
||||
- **空间增强(需同步标签变换)**:水平翻转时自动交换左右舵机角度(s1↔s2, s3↔s4, s5↔s6, s7↔s8),可将训练数据量翻倍。需要了解舵机配对关系后实现。
|
||||
- **动作噪声增强**:`pretrain_mt.py` 已支持 `--noise-std`,`train.py` 可同样引入,对 action_history 和 target 施加高斯扰动。
|
||||
- **更多 `.mt` 脚本**:从根本上扩大动作多样性,是提升模型泛化能力的最有效手段。
|
||||
@@ -4,8 +4,7 @@
|
||||
> 模块:actpred
|
||||
> 作者:自动生成
|
||||
> 日期:2026-08-15
|
||||
> 前置:`actpred_zscore_baseline_20260815.md`(模型 B 基线)、
|
||||
> `actpred_onehot_multihead_dev_plan_20260814.md`(one-hot 方案)
|
||||
> 前置:`actpred_zscore_baseline_20260815.md`(模型 B 基线)
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -5,8 +5,7 @@
|
||||
> 作者:自动生成
|
||||
> 日期:2026-08-15
|
||||
> 相关实验:`pybot_actpred_dataset_training_eval_20260814.md`(A/B 对照)、
|
||||
> `actpred_fixed_range_weighted_experiment_20260814.md`(C/D/E 失败系列)、
|
||||
> `actpred_onehot_multihead_dev_plan_20260814.md`(F 系列结构性方案)
|
||||
> `actpred_fixed_range_weighted_experiment_20260814.md`(C/D/E 失败系列)
|
||||
|
||||
---
|
||||
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -1,358 +0,0 @@
|
||||
# aibox CLI 单次调用模式 · 开发文档
|
||||
|
||||
> 文档类型:功能开发计划
|
||||
> 模块:aibox(单 AI 智能体系统)
|
||||
> 日期:2026-08-17
|
||||
> 状态:方案研究(待评审)
|
||||
> 影响范围:`aibox/aibox.py`(仅增量改动,WebSocket 模式零回归)
|
||||
|
||||
---
|
||||
|
||||
## 1. 概述
|
||||
|
||||
### 1.1 背景与动机
|
||||
|
||||
aibox 目前只有两种使用方式:
|
||||
|
||||
| 使用方式 | 入口 | 特点 |
|
||||
|----------|------|------|
|
||||
| WebSocket 服务器 | `python aibox.py` → `ws://localhost:8765` | 长连接、多轮对话、前端/BotOps 交互(核心场景) |
|
||||
| 备忘录 CLI | `--check-memo` / `--memo-add` / `--memo-list` / `--memo-complete` | 只做任务 CRUD,不触发 AI 对话 |
|
||||
|
||||
缺少一种**"问一次就退"**的 CLI 一次性对话模式。这导致以下场景无法覆盖:
|
||||
|
||||
1. **脚本化调用**:无法在 shell 脚本 / cron / CI 流水线中把 AI 当命令行工具用(如 `aibox --once "统计当前目录 py 文件行数" | grep ...`)
|
||||
2. **资源受限环境**:树莓派 / NPU 盒子上不想常驻一个 WebSocket 服务进程,只想按需拉起、用完即退
|
||||
3. **调试与测试**:复现某个 prompt 行为必须启动服务器 + 前端连接,链路太重
|
||||
4. **进程级隔离执行**:Worker 子代理当前是**同进程线程**(`run_full_subagent` → `DeepSeekChat`),子代理崩溃/死循环会影响主服务;有了 CLI 单次模式后,可演进为"独立子进程执行",OS 级隔离
|
||||
|
||||
### 1.2 设计目标
|
||||
|
||||
1. **新增 `--once QUERY` 参数**:执行一次性 AI 对话(含完整工具调用链),结果流式输出到 stdout,完成后进程退出
|
||||
2. **纯增量、零回归**:不改动 WebSocket 模式任何现有行为;不加参数时行为与现状完全一致
|
||||
3. **复用现有架构**:复用 `EnhancedAIManager` / `DeepSeekChat` / Orchestrator-Worker 工具分层 / 三层中断机制,不重写对话逻辑
|
||||
4. **管道友好**:stdout 只出 AI 回答正文,过程行(🔧 工具调用等)走 stderr,退出码语义化,可被 shell 脚本安全组合
|
||||
5. **状态隔离明确**:单次模式默认不加载历史对话、不写对话归档,记忆写入显式可控
|
||||
|
||||
### 1.3 非目标
|
||||
|
||||
- 不做多轮对话 CLI(交互式 REPL 已由 `client.py` 覆盖)
|
||||
- 不做 HTTP/REST 单次接口(如需可在后续批次评估,复用同一内核)
|
||||
- 不改变 Worker 子代理的线程内执行方式(进程化是后续可选演进,不在本期范围)
|
||||
|
||||
---
|
||||
|
||||
## 2. 现状分析
|
||||
|
||||
### 2.1 入口结构(aibox.py)
|
||||
|
||||
```
|
||||
parse_args() # L2786: argparse 定义
|
||||
main() # L2927: 主入口
|
||||
├── 语言/配置/调试覆盖 # L2931-2959
|
||||
├── 目录创建 + 日志初始化 # L2961-2973
|
||||
├── SubprocessCommandExecutor # L2976
|
||||
├── EnhancedAIManager # L2978-2981
|
||||
├── handle_check_commands() # L2983: --check-memo 处理后 return
|
||||
├── handle_memo_commands() # L2986: --memo-* 处理后 return
|
||||
└── manager.run() # L2990: WebSocket 服务器(唯一终点)
|
||||
```
|
||||
|
||||
关键观察:
|
||||
|
||||
- `handle_check_commands` / `handle_memo_commands` 已经是**"处理后即退出"的 CLI 模式先例**,单次对话模式可完全沿用该模式:在 `manager.run()` 之前插入 `handle_once_mode(manager)`,命中则执行并 return
|
||||
- `EnhancedAIManager.add_ai()`(L2584)已封装完整的 AI 实例组装(memory_db / memo_db / history_manager / prompt_manager / command_executor 注入),单次模式可直接复用
|
||||
- `DeepSeekChat.think_and_respond(input_text, output_callback)`(L1308)是同步阻塞函数,天然适合 CLI 直接调用——不需要 asyncio 事件循环
|
||||
|
||||
### 2.2 output_callback 消息类型(流式输出契约)
|
||||
|
||||
`think_and_respond` 通过 `output_callback(msg_type, content)` 实时推送:
|
||||
|
||||
| msg_type | 语义 | CLI 模式去向 |
|
||||
|----------|------|--------------|
|
||||
| `chunk` | AI 回答增量片段(**权威内容来源**) | **stdout**(流式 write + flush) |
|
||||
| `complete` | 收尾信号(内容与 chunk 累积重复) | 仅作收尾标记,不重复输出 |
|
||||
| `reasoning_start/reasoning/reasoning_end` | 思考模式流式 | stderr(可用 `--show-reasoning` 控制) |
|
||||
| `line` | 过程行:🔧 工具调用、工具输出、💬 前缀 | stderr |
|
||||
| `ai_break` | 工具调用分段信号 | 忽略(CLI 无 widget 概念) |
|
||||
| `error` | API/解析错误 | stderr + 退出码 2 |
|
||||
| `dialogue_ended` | AI 主动结束对话 | 正常收尾 |
|
||||
|
||||
### 2.3 需要处理的技术要点
|
||||
|
||||
| 要点 | 现状 | 风险与对策 |
|
||||
|------|------|-----------|
|
||||
| **压缩输出污染** | `_think_and_respond_impl` L1346 中 `print("📦 对话压缩...")` 直接打 stdout | 单次模式默认不加载历史 → 消息极少 → 几乎不触发压缩;另需确认该 print 走 stderr 或抑制(Batch C 处理) |
|
||||
| **记忆写入并发** | 多个 CLI 进程并发写 `memories.db`(SQLite) | 默认 `--once` 模式 AI 侧仍可调用 `memory` 工具(orchestrator profile);SQLite 有文件锁兜底,低频并发无碍;提供 `--no-save-memory` 开关时通过提示词引导(见 §3.4) |
|
||||
| **Ctrl+C 中断** | 三层取消机制(Event + HTTP close + 子进程 kill)依赖 `cancel()` 被调用 | CLI 模式注册 `signal.SIGINT` handler 调 `manager.cancel()`,与 WebSocket `handle_interrupt` 路径等价 |
|
||||
| **子代理输出转发** | `SubAgentTool._subagent_output` 经 `send_line` → `output_queue` → 回调链 | CLI 模式的 `output_callback` 会自然收到子代理的 `line`/`reasoning`,全部导流 stderr,无需特殊处理 |
|
||||
| **对话归档写入** | `global_exit_handler` 会调 `history_manager.end_current_conversation()` | 单次模式传 `exit_callback=None`(与 TaskScheduler 的任务实例一致),AI 主动 `save_memory_and_end_conversation` 时不触发全局归档 |
|
||||
|
||||
---
|
||||
|
||||
## 3. 设计方案
|
||||
|
||||
### 3.1 命令行接口
|
||||
|
||||
```bash
|
||||
# 基本用法:一次性对话,结果输出到 stdout
|
||||
python aibox.py --once "帮我统计当前目录下所有 py 文件的行数"
|
||||
|
||||
# 静默模式:只输出最终回答正文(工具过程完全静默,适合管道)
|
||||
python aibox.py --once "把 README.md 翻译成英文摘要" --quiet
|
||||
|
||||
# 显式控制
|
||||
python aibox.py --once "..." --no-save-memory # 不写入记忆数据库
|
||||
python aibox.py --once "..." --timeout 120 # 全局超时 120 秒(默认 300)
|
||||
python aibox.py --once "..." --show-reasoning # stderr 显示思考过程(默认不显示)
|
||||
```
|
||||
|
||||
新增参数(全部可选,`--once` 之外均有默认值):
|
||||
|
||||
| 参数 | 类型 | 默认 | 说明 |
|
||||
|------|------|------|------|
|
||||
| `--once QUERY` | str | — | 一次性对话模式:执行 QUERY 后退出,不启动 WebSocket |
|
||||
| `--quiet` | flag | false | 仅 stdout 输出回答正文;过程行/thinking 全静默 |
|
||||
| `--no-save-memory` | flag | false | 禁止本次对话写入记忆数据库 |
|
||||
| `--timeout SEC` | int | 300 | 全局执行超时,超时 kill 并退出码 4 |
|
||||
| `--show-reasoning` | flag | false | 在 stderr 显示思考过程流 |
|
||||
|
||||
### 3.2 退出码语义
|
||||
|
||||
| 退出码 | 含义 |
|
||||
|--------|------|
|
||||
| 0 | 成功(AI 正常完成回答) |
|
||||
| 1 | 工具执行链路异常(如 run_command 失败且 AI 报告失败) |
|
||||
| 2 | API 请求失败(网络/HTTP/JSON 解析,服务端已回滚消息) |
|
||||
| 3 | 用户中断(Ctrl+C) |
|
||||
| 4 | 全局超时(--timeout 到期) |
|
||||
|
||||
> Batch A 只保证 0 / 2 / 3 / 4;退出码 1 的判定(从 AI 文本中识别失败)在 Batch B 实现,初版可统一归 0。
|
||||
|
||||
### 3.3 核心实现(Batch A,约 60~80 行增量)
|
||||
|
||||
**改动点 1:`parse_args()`(aibox.py L2786)** 追加 §3.1 的 5 个参数定义。
|
||||
|
||||
**改动点 2:新增 `handle_once_mode()`**,插入 `main()` 中 `handle_memo_commands()` 之后、`manager.run()` 之前:
|
||||
|
||||
```python
|
||||
def handle_once_mode(manager, args) -> bool:
|
||||
"""一次性对话模式:命中 --once 则执行并返回 True(进程退出)。"""
|
||||
if not args.once:
|
||||
return False
|
||||
|
||||
ai = manager.add_ai()
|
||||
ai.history_manager = None # 单次模式不加载/不写对话历史
|
||||
if args.no_save_memory:
|
||||
ai.memory_db = None # 工具注册时已注入;此处断引用防止运行期写入
|
||||
# 注:Batch A 中 memory 工具已在 _register_tools 注册;
|
||||
# 完整禁用见 §3.4 说明
|
||||
|
||||
exit_code = [0]
|
||||
|
||||
def cli_callback(msg_type, content):
|
||||
if msg_type == "chunk":
|
||||
sys.stdout.write(content)
|
||||
sys.stdout.flush()
|
||||
elif msg_type == "complete":
|
||||
sys.stdout.write("\n")
|
||||
sys.stdout.flush()
|
||||
elif msg_type == "error":
|
||||
print(content, file=sys.stderr)
|
||||
exit_code[0] = 2
|
||||
elif msg_type == "line" and not args.quiet:
|
||||
sys.stderr.write(content)
|
||||
sys.stderr.flush()
|
||||
elif msg_type.startswith("reasoning") and args.show_reasoning:
|
||||
sys.stderr.write(content)
|
||||
sys.stderr.flush()
|
||||
# ai_break / dialogue_ended 忽略
|
||||
|
||||
# Ctrl+C → 三层取消(与 WebSocket interrupt 路径等价)
|
||||
def _sigint(sig, frame):
|
||||
manager.cancel()
|
||||
signal.signal(signal.SIGINT, _sigint)
|
||||
|
||||
# 全局超时:守护线程到点 cancel
|
||||
timer = threading.Timer(args.timeout, manager.cancel)
|
||||
timer.daemon = True
|
||||
timer.start()
|
||||
try:
|
||||
response = ai.think_and_respond(args.once, cli_callback)
|
||||
if response is None or (isinstance(response, str)
|
||||
and response.startswith(f"\n{i18n.get('error_prefix')}")):
|
||||
exit_code[0] = max(exit_code[0], 2)
|
||||
# 被 cancel 打断时 think_and_respond 返回 interrupt 文案
|
||||
if ai._is_cancelled() and timer.is_alive() is False:
|
||||
exit_code[0] = 4
|
||||
finally:
|
||||
timer.cancel()
|
||||
ai.cancel() # 兜底清理子进程
|
||||
|
||||
sys.exit(exit_code[0])
|
||||
```
|
||||
|
||||
**改动点 3:`main()` 一行插入**:
|
||||
|
||||
```python
|
||||
if handle_memo_commands(manager):
|
||||
return
|
||||
|
||||
if handle_once_mode(manager, args): # 新增
|
||||
return
|
||||
|
||||
manager.run()
|
||||
```
|
||||
|
||||
> 注:`handle_once_mode` 内部 `sys.exit()`,return 分支实际不可达,保留 return 仅为与 `handle_check_commands` / `handle_memo_commands` 模式一致。
|
||||
|
||||
### 3.4 记忆写入的处理
|
||||
|
||||
`--no-save-memory` 的**完整实现**需要在工具注册层禁用 `MemoryTool`,但 `_register_tools` 依赖构造期注入的 `memory_db`。Batch A 采用轻量方案:
|
||||
|
||||
- **Batch A**:仅断 `ai.memory_db = None` 引用 + 文档声明"`--once` 场景不建议依赖记忆";AI 若调用 `memory` 工具会因注入的 db 为 None 由工具自身报错兜底(`MemoryTool.execute` 内部访问 db 返回错误字符串,不会崩溃)
|
||||
- **Batch B**:`DeepSeekChat.__init__` 增加 `enable_memory: bool = True` 形参,为 False 时 `_register_tools` 跳过 `MemoryTool` 注册,工具 schema 不出现在请求中,LLM 无从调用——这才是干净方案
|
||||
|
||||
### 3.5 与既有架构约束的对齐
|
||||
|
||||
| 项目硬约束 | 本方案如何满足 |
|
||||
|-----------|---------------|
|
||||
| WebSocket localhost:8765 通信 | 不受影响;`--once` 不启动 WebSocket,互斥使用 |
|
||||
| Worker 仅 7 种执行工具 | 不受影响;单次模式沿用 orchestrator profile,工具分层不变 |
|
||||
| 工具输出 ≤1MB + 超时中断 | 不受影响;复用 `SubprocessCommandExecutor` |
|
||||
| 单向依赖 common←subsystems←tools←aibox | 满足;改动全部在 aibox.py 聚合层 |
|
||||
| 新增/变更模块三处同步(根README/docs/MODULES/ARCHITECTURE) | 本功能非新模块,但落地时需同步更新 `aibox/README.md` 的"命令行参数"章节 |
|
||||
|
||||
---
|
||||
|
||||
## 4. 分批次实施计划
|
||||
|
||||
遵循项目"增量重构、低风险先行"惯例:
|
||||
|
||||
### Batch A:最小可用版(核心,建议先做)
|
||||
|
||||
- [ ] `parse_args()` 新增 `--once` / `--quiet` / `--timeout` 三个参数
|
||||
- [ ] 新增 `handle_once_mode()`(§3.3 实现)
|
||||
- [ ] `main()` 插入调用点
|
||||
- [ ] SIGINT handler + 全局超时 Timer
|
||||
- [ ] 退出码 0/2/3/4
|
||||
- [ ] **验收**:`python aibox.py --once "现在几点了"` 正常输出并退出;`echo $?` 为 0;WebSocket 模式回归(启动服务器 + client.py 对话一轮)
|
||||
|
||||
### Batch B:体验与边界补齐
|
||||
|
||||
- [ ] `--no-save-memory` / `--show-reasoning` 参数落地
|
||||
- [ ] `DeepSeekChat.__init__` 增加 `enable_memory` 开关,干净禁用 MemoryTool
|
||||
- [ ] `_think_and_respond_impl` L1346 压缩提示 `print` 改走 stderr(或注入 print 回调)
|
||||
- [ ] 退出码 1(工具失败识别)
|
||||
- [ ] 补充 `tests/test_once_mode.py`:mock LLM 客户端,验证 stdout/stderr 分流、退出码、超时路径
|
||||
- [ ] 更新 `aibox/README.md` 命令行参数章节
|
||||
|
||||
### Batch C:安全加固(可选,按需)
|
||||
|
||||
- [ ] SQLite WAL 模式(`PRAGMA journal_mode=WAL`),缓解多 CLI 进程并发写锁等待
|
||||
- [ ] `--worker-only` 开关:`tool_profile=worker` 纯执行面模式,供其他模块以子进程方式调用 AI 执行器(为"Worker 进程级隔离"演进铺路)
|
||||
- [ ] `--model` / `--temperature` 覆盖参数
|
||||
|
||||
### 明确不做(本期)
|
||||
|
||||
- 多轮交互式 CLI REPL(`client.py` 已覆盖)
|
||||
- HTTP 单次接口
|
||||
- Worker 子进程化(架构演进议题,单独立项)
|
||||
|
||||
---
|
||||
|
||||
## 5. 风险清单
|
||||
|
||||
| 风险 | 等级 | 缓解 |
|
||||
|------|------|------|
|
||||
| 冷启动开销(PromptManager 加载 + SQLite 初始化 + LLM 握手 ≈1~3s) | 低 | 可接受取舍;批量场景文档引导用 WebSocket 长连接 |
|
||||
| CLI 模式与 WebSocket 模式代码 drift | 中 | 复用 `EnhancedAIManager.add_ai()`,单次模式新增代码控制在 80 行内,核心对话逻辑零复制 |
|
||||
| `print` 污染 stdout 破坏管道 | 中 | Batch A 用 `--quiet` 兜底;Batch B 根治压缩提示 print |
|
||||
| 并发进程 SQLite 写锁 | 低 | 低频无碍;Batch C WAL 加固 |
|
||||
| Ctrl+C 时子进程泄漏 | 低 | 复用现成三层取消;`finally` 中 `ai.cancel()` 兜底 |
|
||||
|
||||
---
|
||||
|
||||
## 6. 测试方案
|
||||
|
||||
### 6.1 单元测试(tests/test_once_mode.py,Batch B)
|
||||
|
||||
| 用例 | 验证点 |
|
||||
|------|--------|
|
||||
| 正常问答 | mock LLM 返回纯文本 → stdout 含回答、退出码 0 |
|
||||
| 工具调用链 | mock LLM 返回 tool_calls → line 走 stderr、chunk 走 stdout |
|
||||
| API 失败 | mock LLM 抛 RequestException → 退出码 2、stderr 含错误 |
|
||||
| 超时 | mock LLM 阻塞 → Timer 触发 cancel → 退出码 4 |
|
||||
| SIGINT | 发送 SIGINT → 三层取消触发 → 退出码 3 |
|
||||
| `--quiet` | 过程行不输出到任何流 |
|
||||
| `--no-save-memory` | memories.db 无新增行 |
|
||||
|
||||
### 6.2 回归测试
|
||||
|
||||
- [ ] `py -m pytest aibox/tests -q` 现有 232 项全过
|
||||
- [ ] WebSocket 模式手测:启动服务器 + `client.py` 完成一轮对话 + Ctrl+C 中断
|
||||
|
||||
### 6.3 手动验收脚本
|
||||
|
||||
```bash
|
||||
# 基础问答(无工具)
|
||||
python aibox.py --once "用一句话介绍你自己" ; echo "exit=$?"
|
||||
|
||||
# 工具调用链(会派生 worker 子代理执行命令)
|
||||
python aibox.py --once "在 ~/.aibox/workspace 下创建一个 hello.txt 写入 hello" ; echo "exit=$?"
|
||||
|
||||
# 管道组合
|
||||
python aibox.py --once --quiet "列出当前目录 py 文件" | grep ".py" | wc -l
|
||||
|
||||
# 超时路径
|
||||
python aibox.py --once --timeout 5 "运行 sleep 60" ; echo "exit=$?" # 期望 4
|
||||
```
|
||||
|
||||
## 7. 实施状态(2026-08-17 落地记录)
|
||||
|
||||
> 本节为实施后回填,与第 3/4 节计划存在少量偏差,以本节为准。
|
||||
|
||||
### 7.1 提交与批次
|
||||
|
||||
| 批次 | 提交 | 内容 |
|
||||
|---|---|---|
|
||||
| Batch A | `ce3fb7f` | `--once` / `--quiet` / `--once-timeout`(300s 默认)+ `handle_once_mode` + SIGINT/超时复用三层取消 |
|
||||
| Batch B | `9e0d39c` | `--no-save-memory`(`enable_memory` 构造开关干净禁用 MemoryTool)、`--show-reasoning`(思考流默认静默走 stderr)、压缩提示改走 output_callback、退出码 1(orchestrator 层工具异常)+ `test_once_mode.py` 10 项 |
|
||||
| Batch C | `a5af4bb` | SQLite WAL、`--worker-only`(worker profile + 精简 system prompt 防人格污染)、`--model` / `--temperature`、子代理失败标记透传(`_SUBAGENT_TOOL_FAILED_MARKER` → orchestrator `stats.tool_failed`)+ 5 项测试 |
|
||||
| Batch D | `d03dbc7` | 多模型档案 `models` 配置节(本档案为实施中新增需求,超出原计划范围)+ `test_models_profile.py` 17 项 + README 更新 |
|
||||
|
||||
### 7.2 与计划的偏差
|
||||
|
||||
1. **超时参数名**:计划 `--timeout`,落地为 `--once-timeout`(避免与 `api.timeout` 语义混淆)
|
||||
2. **`models` 默认留空 `{}`**:计划未涉及;实施中发现默认档案硬编码 `DEEPSEEK_API_KEY` / 官方 URL 会覆盖用户自定义配置(本机为 `ARK_API_KEY` + 火山方舟端点,曾导致 401),故默认空 dict,未配置档案时零破坏回退旧 `api` 扁平节
|
||||
3. **`--worker-only` 需要精简 system prompt**:原计划仅切 tool_profile;实测 orchestrator 的人格/身份提示词会污染 worker 语义(自称"老板"、主动写文件、多轮不收敛),已对齐 `run_full_subagent` 的子代理语义修复
|
||||
4. **退出码 1 的边界**:仅覆盖 orchestrator 直接感知的工具异常;worker 子代理内的命令失败经 Batch C 的标记透传后同样计入。子命令失败但 AI 正常汇报时退出码为 0(AI 视角任务完成)
|
||||
|
||||
### 7.3 验证结果
|
||||
|
||||
- 单测 **264 项全过**(once 模式 15 + 多模型档案 17 新增)
|
||||
- 冒烟:基础问答 exit 0 / 工具调用链(子代理执行 echo)exit 0 / 超时路径 exit 4 / `--no-save-memory --quiet` stdout 干净 / `--show-reasoning` 思考流走 stderr / `--temperature` 透传 / 不存在的模型名 exit 2
|
||||
|
||||
### 7.4 明确不做(维持原判)
|
||||
|
||||
WebSocket 会话恢复、`--once` 结果缓存、批处理多 query——见第 4 节"明确不做"。
|
||||
|
||||
---
|
||||
|
||||
## 8. 后续演进方向(不在本期)
|
||||
|
||||
1. **Worker 进程级隔离**:`run_full_subagent` 的线程内 `DeepSeekChat` → `subprocess.run(["python", "aibox.py", "--once", "--worker-only", ...])`,子代理崩溃不波及主服务,可用 OS 用户/chroot 做权限隔离
|
||||
2. **HTTP 单次接口**:`POST /chat` 无状态接口,与 `--once` 共享内核
|
||||
3. **结构化输出**:`--json` 输出 `{answer, tool_calls, exit_reason}` JSON,供程序消费
|
||||
|
||||
---
|
||||
|
||||
## 9. 决策记录
|
||||
|
||||
| 决策点 | 结论 | 理由 |
|
||||
|--------|------|------|
|
||||
| 新入口文件还是复用 aibox.py | 复用 aibox.py + `--once` 参数 | 沿用 `--check-memo` 先例;避免维护两份入口 drift |
|
||||
| stdout 输出协议 | chunk 流式 + 末尾换行 | 与 WebSocket 的 chunk 权威语义一致;管道可逐段消费 |
|
||||
| 过程行去向 | stderr | Unix 管道惯例:stdout 是数据、stderr 是元信息 |
|
||||
| 单次模式是否加载历史 | 否 | "一次性"语义明确;避免历史污染导致行为不可复现 |
|
||||
| 是否默认写记忆 | 默认与 orchestrator 一致(可写),`--no-save-memory` 显式关闭 | 与 WebSocket 模式行为对齐,减少心智负担 |
|
||||
@@ -1,183 +0,0 @@
|
||||
# aibox 文件操作与网络操作命令行化 开发计划
|
||||
|
||||
## 目标
|
||||
|
||||
将 aibox 中所有 Python 原生实现的文件操作和网络操作改为通过系统命令行工具实现,统一工具执行模式,减少 Python 依赖。
|
||||
|
||||
## 涉及模块
|
||||
|
||||
| 现有类 | 作用 | 替换方案 |
|
||||
|--------|------|----------|
|
||||
| `FileToolBase` | 文件工具基类 | 简化为命令行工具基类 |
|
||||
| `ReadFileTool` | 读文件 | `cat` / `head` / `tail` / `wc` |
|
||||
| `WriteFileTool` | 写文件 | `mkdir -p` + shell 重定向 |
|
||||
| `ListFilesTool` | 列目录 | `ls` / `find` / `du` |
|
||||
| `SimpleWebClient` | HTTP 客户端 | `curl` |
|
||||
| `WebPage` | 网页解析对象 | 移除,curl + grep/sed 替代 |
|
||||
| `SimpleWebCrawlerTool` | 网页爬虫工具 | `curl` + 文本处理管道 |
|
||||
|
||||
## 阶段一:文件操作命令行化
|
||||
|
||||
> **实现说明(已落地)**:阶段一采用「工具返回 markdown 提示词」的方式而非工具内部调子进程。`read_file` / `write_file` / `list_files` 三个工具的 `execute()` 不再执行 Python 文件 I/O,而是返回一段 markdown 提示词,指导智能体改用 `run_command` 工具执行对应的命令行命令(`head`/`tail`/`sed`/`wc`/`file`/`cut`/`cat`/`ls`/`du`/`find`)。工具接口(名称与参数)保持不变,安全限制(500 行、10KB/行、10MB 文件、路径 `shlex.quote` 转义)在提示词中注明。
|
||||
|
||||
### 1.1 ReadFileTool 重构
|
||||
|
||||
用以下命令行工具替代 Python 文件 I/O:
|
||||
|
||||
| 原功能 | 命令行实现 |
|
||||
|--------|-----------|
|
||||
| `head=True` 读取开头 N 行 | `head -n N` |
|
||||
| `tail=True` 读取末尾 N 行 | `tail -n N` |
|
||||
| 指定行范围 `start_line ~ end_line` | `sed -n 'start,endp'` |
|
||||
| 统计行数 `_count_lines` | `wc -l` |
|
||||
| 文件编码检测 | `file -b --mime-encoding` |
|
||||
| 单行超长截断 | `cut -c 1-10000` |
|
||||
|
||||
保留的安全限制:
|
||||
- 文件大小限制 `MAX_FILE_SIZE`:通过 `wc -c` 或 `stat -f%z` 检查
|
||||
- 行数上限 500 行:`head -n 500` 作为最后一道防线
|
||||
- 单行 10KB 限制:`cut -c 1-10000`
|
||||
|
||||
### 1.2 WriteFileTool 重构
|
||||
|
||||
| 原功能 | 命令行实现 |
|
||||
|--------|-----------|
|
||||
| 创建目录 | `mkdir -p` |
|
||||
| 覆盖写入 `mode=w` | `content > filepath` |
|
||||
| 追加写入 `mode=a` | `content >> filepath` |
|
||||
| 检查文件大小 | `wc -c` 或 `stat -f%z` |
|
||||
| 路径安全检查 | `realpath` / `test -f` |
|
||||
|
||||
### 1.3 ListFilesTool 重构
|
||||
|
||||
| 原功能 | 命令行实现 |
|
||||
|--------|-----------|
|
||||
| 列出目录内容 | `ls -la` 或 `find . -maxdepth 1` |
|
||||
| 隐藏文件过滤 | `ls -la` 默认不显示 `.` 开头的隐藏文件(或 `find` 排除) |
|
||||
| 文件大小统计 | `du -sh *` 或 `ls -lh` |
|
||||
| 目录/文件区分 | `ls -F` 或 `find -type f/d` |
|
||||
|
||||
### 1.4 FileToolBase 简化
|
||||
|
||||
移除 `normalize_path()`、`check_path_safety()` 中 Python `os` 模块调用,改为:
|
||||
- 路径规范化:`realpath`
|
||||
- 文件存在性检查:`test -f` / `test -d`
|
||||
- 文件大小检查:`wc -c` 或 `stat`
|
||||
|
||||
### 1.5 安全边界保留
|
||||
|
||||
- `MAX_FILE_SIZE`(默认 10MB):通过命令检查
|
||||
- 行数上限 500 行:`head -n 500` 兜底
|
||||
- 单行长度上限 10KB:`cut -c 1-10000` 兜底
|
||||
- 所有路径参数 `shlex.quote()` 转义
|
||||
|
||||
## 阶段二:网络操作命令行化
|
||||
|
||||
> **实现说明(已落地)**:阶段二沿用阶段一的方式,`web` 工具(`SimpleWebCrawlerTool`)的 `execute()` 不再使用 `requests`/`BeautifulSoup`,而是返回 markdown 提示词,指导智能体改用 `run_command` 执行 `curl` 命令完成各 action(`get`/`text`/`title`/`links`/`search`)。`WebPage`、`SimpleWebClient` 类已删除;`urljoin`/`urlparse`/`BeautifulSoup` 导入已清理。代理通过 `get_proxy()` 读取配置并拼出 `--socks5` 参数。注意 `requests` 仍在 `aibox.py` 其他位置(API 调用、DeepSeek 请求)使用,故保留其依赖。
|
||||
|
||||
### 2.1 SimpleWebClient → curl
|
||||
|
||||
| 原功能 | 命令行实现 |
|
||||
|--------|-----------|
|
||||
| HTTP GET 请求 | `curl -sL --max-time N` |
|
||||
| 自定义 User-Agent | `curl -A "..."` |
|
||||
| SOCKS5 代理 | `curl --socks5 host:port` 或环境变量 `ALL_PROXY` |
|
||||
| 请求间隔控制 | `SubprocessCommandExecutor` 层面控制 |
|
||||
| 响应状态码 | `curl -sL -o /dev/null -w '%{http_code}'` |
|
||||
| 下载大小限制 | `curl --max-filesize N` 或管道 `head -c N` |
|
||||
|
||||
### 2.2 WebPage 移除
|
||||
|
||||
`WebPage` 类不再需要,其功能由命令行管道实现:
|
||||
|
||||
| 原属性/方法 | 命令行实现 |
|
||||
|------------|-----------|
|
||||
| `title` | `curl -sL | grep -oP '<title>\K[^<]+'` |
|
||||
| `text` | `curl -sL | lynx -dump -stdin` 或 `html2text` |
|
||||
| `links` | `curl -sL | grep -oP 'href="[^"]*"'` |
|
||||
| `status_code` | `curl -sL -o /dev/null -w '%{http_code}'` |
|
||||
| `headers` | `curl -sL -I` |
|
||||
|
||||
### 2.3 SimpleWebCrawlerTool 重构
|
||||
|
||||
| action | 命令行实现 |
|
||||
|--------|-----------|
|
||||
| `get` | `curl -sL` 获取完整页面信息 |
|
||||
| `text` | `curl -sL | lynx -dump -stdin` 或 `html2text` |
|
||||
| `title` | `curl -sL | grep -oP '<title>\K[^<]+'` |
|
||||
| `links` | `curl -sL | grep -oP 'href="[^"]*"'` |
|
||||
| `search` | `curl -sL | grep -c 'keyword'` 或 `grep -q` |
|
||||
|
||||
### 2.4 下载大小限制
|
||||
|
||||
- 响应体:`curl --max-filesize` 或 `curl | head -c 100M` 管道限制
|
||||
- 配置项 `web.max_download_size`(默认 100MB)
|
||||
|
||||
## 阶段三:清理与收尾
|
||||
|
||||
> **实现说明(已落地)**:
|
||||
> - 3.1 `requirements.txt` 已移除 `beautifulsoup4`;**`requests` 保留**(仍在 `aibox.py` 的 DeepSeek API 调用 `ConversationCompressor` 与 `DeepSeekChat` 中使用,属核心功能,无法移除)。
|
||||
> - 3.2 已统一:文件/网页工具改为返回 markdown 提示词,实际 I/O 全部委托给 `run_command`(`SubprocessCommandExecutor`),天然复用其 1MB 输出上限、超时与中断取消。
|
||||
> - 3.3 中断支持已随 `SubprocessCommandExecutor.cancel()` 覆盖。
|
||||
> - 3.4 默认提示词 `3_file_operations.md`、`4_web_access.md` 已更新为命令行走指导(`head`/`sed`/`cat`/`ls`/`curl` 等)。
|
||||
> - 3.5 `README.md` 工具表与安全说明已更新。
|
||||
|
||||
### 3.1 移除 Python 依赖
|
||||
|
||||
从 `requirements.txt` 删除:
|
||||
|
||||
```diff
|
||||
- requests
|
||||
- beautifulsoup4
|
||||
```
|
||||
|
||||
### 3.2 统一使用 SubprocessCommandExecutor
|
||||
|
||||
所有新工具统一通过 `SubprocessCommandExecutor` 执行命令,与 `CommandLineTool` 保持一致:
|
||||
- 输出上限 1MB(超出自动终止进程并截断)
|
||||
- 超时控制
|
||||
- 中断取消(`cancel()` → `process.kill()`)
|
||||
|
||||
### 3.3 中断支持
|
||||
|
||||
- 利用已有的 `SubprocessCommandExecutor.cancel()` 机制
|
||||
- 确保 Ctrl+C 可终止长时间运行的 curl 下载或文件操作
|
||||
|
||||
### 3.4 更新提示词
|
||||
|
||||
`~/.aibox/prompts/` 中相关文件更新:
|
||||
|
||||
| 文件 | 更新内容 |
|
||||
|------|---------|
|
||||
| `3_file_operations.md` | 文件操作描述改为命令行方式 |
|
||||
| `4_web_access.md` | 网络访问描述改为 curl 方式 |
|
||||
|
||||
### 3.5 更新文档
|
||||
|
||||
- `SKILLS_FORMAT.md`:技能开发示例中文件/网络操作更新
|
||||
- `README.md`:更新依赖说明和架构图
|
||||
|
||||
## 不做的事
|
||||
|
||||
- **不改变工具接口**:`read_file`、`write_file`、`list_files`、`web` 四个工具名称和参数签名保持不变
|
||||
- **不改变 `CommandLineTool`(`run_command`)**:已有的命令执行工具不动
|
||||
- **不改变 `SubprocessCommandExecutor`**:已有安全限制(1MB 输出)保持不变
|
||||
- **不改变记忆/备忘录数据库**:`MemoryDatabase`、`MemoDatabase` 的 SQLite 操作保留
|
||||
|
||||
## 收益
|
||||
|
||||
- **减少依赖**:移除 `requests`、`beautifulsoup4` 两个 Python 库,降低供应链风险
|
||||
- **统一执行模式**:所有外部操作(文件、网络、命令)全部走 `SubprocessCommandExecutor`,统一超时控制、输出限制、中断取消
|
||||
- **更强的可审计性**:AI 操作日志中能看到完整的 `curl` / `cat` / `ls` 命令,更透明
|
||||
- **复用已有基础设施**:`CommandLineTool` 已有的安全防护直接复用
|
||||
|
||||
## 风险与注意事项
|
||||
|
||||
| 风险 | 缓解措施 |
|
||||
|------|----------|
|
||||
| `curl` 不存在 | 启动时检测 `curl` 可用性,缺失时尝试 `wget` 降级或报错 |
|
||||
| 跨平台 `sed`/`grep` 差异 | macOS (BSD) 和 Linux (GNU) 行为差异;使用 POSIX 兼容写法,必要时运行时检测 |
|
||||
| 网页解析精度下降 | `grep` 正则提取 HTML 不可靠;优先用 `lynx -dump` 或 `html2text` 替代 BeautifulSoup 解析 |
|
||||
| 命令注入风险 | 所有用户输入参数必须经过 `shlex.quote()` 转义 |
|
||||
| Proxy 配置 | `curl` 通过 `--socks5` 或 `ALL_PROXY` 环境变量传递,与现有 `set_proxy` 工具集成 |
|
||||
| 性能影响 | 文件操作每步 fork 新进程,但 aibox 不是高吞吐服务,IO 瓶颈远大于进程开销 |
|
||||
@@ -1,184 +0,0 @@
|
||||
# aibox 编排者-执行者工具分层实施方案
|
||||
|
||||
> 状态:待实施 · 日期:2026-08-13 · 基线:aibox v2.8(tools 包抽离后,aibox.py 7753 行)
|
||||
> 前置:`14db64f refactor(aibox): 工具基类 Tool/ToolRegistry 抽离至 tools 包`
|
||||
|
||||
## 1. 背景与目标
|
||||
|
||||
当前 aibox 主对话与完整子代理(`run_full_subagent`)注册**同一份全量工具集**(21 个工具),带来三个问题:
|
||||
|
||||
1. **安全隐患**:子代理可调用 `update_soul` / `update_identity` 改写主人格文件,且携带 `create_subagent` 可无限递归派生。
|
||||
2. **上下文膨胀**:主对话 system prompt 携带全部 21 个工具的 schema 与使用说明,token 成本高。
|
||||
3. **职责混杂**:命令执行、记忆管理、情绪表达交织在同一上下文中,长任务易污染主对话。
|
||||
|
||||
**目标架构**(编排者-执行者分层):
|
||||
|
||||
- **主对话(orchestrator)**:只保留情绪管理、记忆管理、任务调度、子代理派生等"管理面"工具,**不直接接触命令行**。
|
||||
- **子代理(worker)**:唯一拥有命令行/文件/网页等"执行面"工具的角色,独立上下文,无记忆/人格/派生能力。
|
||||
- 轻量子代理(`run_light_subagent`,供 `express_emotion` 使用)保持不变。
|
||||
|
||||
## 2. 现状分析(关键代码位置)
|
||||
|
||||
| 位置 | 说明 |
|
||||
|------|------|
|
||||
| `aibox.py:5667` `DeepSeekChat.__init__` | 无工具集参数,构造即调 `_register_tools` |
|
||||
| `aibox.py:5766` `_register_tools()` | 无条件注册全部 21 个内置工具 + `search_skills` |
|
||||
| `aibox.py:5392` `run_full_subagent()` | 创建完整 `DeepSeekChat(system_prompt=sub_role, ...)`,子代理因此获得全量工具(含 `update_soul`/`create_subagent`)|
|
||||
| `aibox.py:5363` `run_light_subagent()` | 无工具,单次 LLM 补全,供 `ExpressEmotionTool`(:3408) 使用 |
|
||||
| `aibox.py:5524` `SubAgentTool` | `create_subagent` 工具,调 `run_full_subagent` |
|
||||
| `~/.aibox/prompts/` | 运行期提示词(7_extension_skills.md 等指导 AI 直接用 `run_command` 执行技能)|
|
||||
|
||||
## 3. 工具分配方案
|
||||
|
||||
### orchestrator(主对话)— 13 个工具
|
||||
|
||||
| 工具 | 保留理由 |
|
||||
|------|---------|
|
||||
| `memory` / `memo` / `save_memory_and_end_conversation` | 记忆与任务管理是主代理核心职责 |
|
||||
| `express_emotion` / `ask_user` | 情绪管理与用户交互只在主对话有意义 |
|
||||
| `create_subagent` | 唯一的执行入口 |
|
||||
| `search_skills` | 只查不执行,用于路由决策(查到技能后派 worker 执行)|
|
||||
| `reload_skills` / `install_skill` / `create_skill` | 技能管理面;后两者为指引型工具,实际文件写入通过 worker 完成 |
|
||||
| `update_soul` / `update_identity` | 人格进化收敛到主代理(同时消除子代理越权隐患)|
|
||||
| `set_proxy` / `get_proxy` | 代理配置影响 worker 的 web/curl,归管理面 |
|
||||
| `get_current_time` | 任务调度/记忆时间戳需要 |
|
||||
|
||||
### worker(子代理)— 7 个工具
|
||||
|
||||
| 工具 | 说明 |
|
||||
|------|------|
|
||||
| `run_command` | 命令行执行(含技能脚本执行)|
|
||||
| `read_file` / `write_file` / `list_files` | 文件操作三件套(指引型)|
|
||||
| `web` | 网页访问指引 |
|
||||
| `calculator` / `get_current_time` | 执行辅助 |
|
||||
|
||||
### 明确不给 worker 的工具
|
||||
|
||||
`memory` / `memo`(记忆隔离)、`update_soul` / `update_identity`(人格保护)、`create_subagent`(断递归)、`express_emotion`、`ask_user`(**worker 不得向用户提问**,疑问写入返回结果)。
|
||||
|
||||
## 4. 改动清单
|
||||
|
||||
### 4.1 `aibox.py` — 工具 profile 参数化(核心改动)
|
||||
|
||||
**a) 模块级常量**(`DeepSeekChat` 类定义前):
|
||||
|
||||
```python
|
||||
# 工具集 profile:orchestrator=主对话管理面,worker=子代理执行面
|
||||
TOOL_PROFILE_ORCHESTRATOR = "orchestrator"
|
||||
TOOL_PROFILE_WORKER = "worker"
|
||||
```
|
||||
|
||||
**b) `DeepSeekChat.__init__`(:5667)** 增加参数 `tool_profile: str = TOOL_PROFILE_ORCHESTRATOR`,保存为 `self.tool_profile`,`_register_tools` 调用处透传。
|
||||
|
||||
**c) `_register_tools()`(:5766)** 按 profile 分支装配:
|
||||
|
||||
```python
|
||||
def _register_tools(self, exit_callback=None, command_executor=None):
|
||||
if self.tool_profile == TOOL_PROFILE_WORKER:
|
||||
builtin_tools = [
|
||||
CommandLineTool(command_executor or SubprocessCommandExecutor(self.logger), self.logger),
|
||||
ReadFileTool(self.logger),
|
||||
WriteFileTool(self.logger),
|
||||
ListFilesTool(self.logger),
|
||||
SimpleWebCrawlerTool(self.logger),
|
||||
CalculatorTool(self.logger),
|
||||
GetCurrentTimeTool(self.logger),
|
||||
]
|
||||
self.tool_registry.register_many(builtin_tools, is_builtin=True)
|
||||
return # worker 不需要技能管理器与 search_skills
|
||||
|
||||
# —— 以下为 orchestrator 分支,即现有装配逻辑删去执行面 7 工具 ——
|
||||
builtin_tools = [GetCurrentTimeTool(self.logger)]
|
||||
# memory/memo/update_soul/update_identity/save_memory/ask_user/
|
||||
# express_emotion/reload_skills/install_skill/create_skill/
|
||||
# set_proxy/get_proxy/SubAgentTool 保持现有条件注册逻辑
|
||||
...
|
||||
```
|
||||
|
||||
注意:orchestrator 分支从现有列表(:5768-5807)中移除 `CalculatorTool`、`CommandLineTool`、`ReadFileTool`、`WriteFileTool`、`ListFilesTool`、`SimpleWebCrawlerTool`,其余不变。
|
||||
|
||||
**d) `run_full_subagent()`(:5392)** 创建子实例时传 `tool_profile=TOOL_PROFILE_WORKER`:
|
||||
|
||||
```python
|
||||
sub_chat = DeepSeekChat(
|
||||
system_prompt=sub_role,
|
||||
tool_profile=TOOL_PROFILE_WORKER, # 新增
|
||||
command_executor=SubprocessCommandExecutor(),
|
||||
...
|
||||
)
|
||||
```
|
||||
|
||||
### 4.2 子代理角色提示词补强(`run_full_subagent` 内 `sub_role` 拼接)
|
||||
|
||||
现有结尾指令("任务完成后直接给出最终回答")追加 worker 行为约束:
|
||||
|
||||
```
|
||||
【约束】你没有记忆/提问/派生能力。遇到不确定的信息,用工具自行查证;
|
||||
确实无法完成时,在最终回答中明确说明原因与已完成的部分。
|
||||
产物文件一律写入 ~/.aibox/workspace/ 并在回答中给出绝对路径。
|
||||
```
|
||||
|
||||
### 4.3 运行期提示词调整(`~/.aibox/prompts/`,部署后手工更新)
|
||||
|
||||
- `99_SOUL.md`:角色定位改为"编排者"——拆解任务、派生子代理执行、汇总结果;明确自己**不直接执行命令**。
|
||||
- `7_extension_skills.md`:技能执行路径从"直接 run_command"改为"`search_skills` 查到技能 → `create_subagent` 派 worker 执行 → 汇总结论"。
|
||||
- `6_command_execution.md`:改为说明命令执行须经子代理(或直接停用该提示词)。
|
||||
- `10_task_handling.md`:定时任务执行同样走子代理路径。
|
||||
|
||||
> 提示词为运行期数据不在仓库内,实施时在部署机 `~/.aibox/prompts/` 修改,并用 `/reload_prompts` 生效。
|
||||
|
||||
### 4.4 文档更新
|
||||
|
||||
- `aibox/README.md`:工具表按 orchestrator/worker 重新分组;"架构流程"图补子代理层;更新日志加 v2.9 条目。
|
||||
|
||||
## 5. 已知限制与风险
|
||||
|
||||
| 风险 | 说明 | 对策 |
|
||||
|------|------|------|
|
||||
| 延迟与成本上升 | 每次执行多一跳 LLM + 结束摘要(`_summarize_subagent_conversation`) | 简单问答主代理直接回答不派生;提示词明确"仅在需要执行时派生" |
|
||||
| 信息失真 | 主代理只见摘要(8000 字符截断) | worker 产物落盘 workspace,主代理用 `search_skills`/派生复查;后续可提升摘要上限 |
|
||||
| 中断不彻底 | worker 自建 `SubprocessCommandExecutor`(:5392 内),主代理 Ctrl+C 无法 kill worker 的子进程 | 现状已有该缺陷;后续改进:worker executor 注册到主 `_command_executor` 链或全局进程表 |
|
||||
| 任务调度器 | TaskScheduler 触发的对话为主代理,执行类任务须学会派生 | 提示词 10_task_handling.md 明确路径;上线后观察任务执行日志 |
|
||||
| 回退 | 行为变更较大 | 单提交实现,`git revert` 即可完整回退 |
|
||||
|
||||
## 6. 验证方案
|
||||
|
||||
```bash
|
||||
cd /opt/blackbird/aibox
|
||||
|
||||
# 1. 导入冒烟(系统 python 缺 websockets 时用桩模块)
|
||||
python3 -m py_compile aibox.py
|
||||
|
||||
# 2. profile 装配断言(桩掉 websockets/requests 后)
|
||||
python3 - <<'EOF'
|
||||
import sys, types
|
||||
for n in ("websockets",): sys.modules.setdefault(n, types.ModuleType(n))
|
||||
import aibox
|
||||
o = aibox.DeepSeekChat(tool_profile="orchestrator")
|
||||
w = aibox.DeepSeekChat(tool_profile="worker")
|
||||
on = sorted(t.get_name() for t in o.tool_registry.get_all_tools())
|
||||
wn = sorted(t.get_name() for t in w.tool_registry.get_all_tools())
|
||||
assert "run_command" not in on and "create_subagent" in on
|
||||
assert set(wn) == {"run_command","read_file","write_file","list_files",
|
||||
"web","calculator","get_current_time"}
|
||||
assert "update_soul" not in wn and "create_subagent" not in wn
|
||||
print("orchestrator:", on); print("worker:", wn)
|
||||
EOF
|
||||
```
|
||||
|
||||
(注:构造 `DeepSeekChat` 需 `DEEPSEEK_API_KEY`,`_setenv.sh` 已从 `~/.cert/deepseek.key` 导出。)
|
||||
|
||||
**端到端**(部署机):
|
||||
|
||||
1. `bash _run.sh` 启动,BotOps 连接。
|
||||
2. 问"帮我查看系统内存"→ 预期:主代理调 `create_subagent`,任务侧栏见 worker 执行 `free -h`,主代理汇总回答;主对话流中不出现 `run_command` 工具调用。
|
||||
3. 问"记住我喜欢喝绿茶"→ 预期:主代理直接调 `memory`,不派生子代理。
|
||||
4. 问"使用 nmapskills 技能扫描本机"→ 预期:`search_skills` → `create_subagent` 执行技能脚本。
|
||||
5. Ctrl+C 中断一个执行中的子代理任务,确认 `interrupt_ack` 与回滚正常。
|
||||
6. `/tasks` 添加即时任务"统计 ~/.aibox/workspace 文件数",确认调度器路径走子代理执行。
|
||||
|
||||
## 7. 后续可选扩展(不在本期)
|
||||
|
||||
- `config.json` 增加 `agent.worker_model`,worker 用更廉价模型降本。
|
||||
- worker 进程注册到全局取消链,实现彻底中断。
|
||||
- 摘要失真问题严重时可改为 worker 结果直通主对话(`task_process` 已有广播通道)。
|
||||
@@ -1,167 +0,0 @@
|
||||
# aibox 模块化重构进度报告
|
||||
|
||||
> 状态:A/B 批工具 + Wave 1 子系统已抽离 · 日期:2026-08-14 · 基线:aibox v3.0
|
||||
> 测试:171 项全过(10.79s)
|
||||
> 关联文档:
|
||||
> - [aibox_orchestrator_worker_plan.md](aibox_orchestrator_worker_plan.md)(v2.9 编排者-执行者分层方案)
|
||||
> - [aibox_open_dev_plan.md](aibox_open_dev_plan.md)(开放版安全裁剪方案)
|
||||
> - [aibox_SECURITY_AUDIT_20260801.md](../security/aibox_SECURITY_AUDIT_20260801.md)(安全审计基线)
|
||||
|
||||
## 1. 重构目标
|
||||
|
||||
将单文件 8450 行的 `aibox.py` 按职责拆分为多个内聚包,建立清晰的依赖方向与测试基线,降低维护成本、为后续工具/子系统扩展提供稳定骨架。
|
||||
|
||||
**目标架构**(单向依赖,禁止反向引用):
|
||||
|
||||
```
|
||||
common ← subsystems ← tools ← aibox.py
|
||||
```
|
||||
|
||||
| 层 | 职责 | 依赖 |
|
||||
|---|---|---|
|
||||
| `common/` | 配置、常量、i18n、日志 | 标准库 |
|
||||
| `subsystems/` | 命令执行器、提示词、数据库、技能管理 | `common` |
|
||||
| `tools/` | Tool 基类 + 17 个工具类 | `common` + `subsystems` |
|
||||
| `aibox.py` | 核心类聚合(DeepSeekChat/WebSocketHandler/EnhancedAIManager + 调度/压缩/历史) | 上述三包 |
|
||||
|
||||
## 2. 已完成工作
|
||||
|
||||
### 2.1 A 批基础工具抽离(`tools/basic_tools.py`,342 行)
|
||||
|
||||
7 个无外部子系统依赖的工具类:
|
||||
|
||||
| 类 | 工具名 | 说明 |
|
||||
|---|---|---|
|
||||
| `FileToolBase` | — | 文件操作基类(normalize_path / check_path_safety) |
|
||||
| `ReadFileTool` | `read_file` | 读文件指引(head/tail/sed) |
|
||||
| `WriteFileTool` | `write_file` | 写文件指引(heredoc) |
|
||||
| `ListFilesTool` | `list_files` | 列目录指引(ls) |
|
||||
| `SimpleWebCrawlerTool` | `web` | 网页访问指引(curl + 代理) |
|
||||
| `GetCurrentTimeTool` | `get_current_time` | 时间获取 |
|
||||
| `CalculatorTool` | `calculator` | 四则运算 |
|
||||
|
||||
**测试**:[tests/test_basic_tools.py](../../aibox/tests/test_basic_tools.py) 43 项全过
|
||||
|
||||
### 2.2 B 批管理面工具抽离(`tools/` 目录,共 13 类)
|
||||
|
||||
| 文件 | 类 | 工具名 |
|
||||
|---|---|---|
|
||||
| `command_tool.py` | `CommandLineTool` | `run_command` |
|
||||
| `skill_tools.py` | `SearchSkillsTool` / `ReloadSkillsTool` / `InstallSkillTool` / `CreateSkillTool` | `search_skills` / `reload_skills` / `install_skill` / `create_skill` |
|
||||
| `db_tools.py` | `MemoryTool` / `MemoTool` | `memory` / `memo` |
|
||||
| `prompt_tools.py` | `UpdateIdentityTool` / `UpdateSoulTool` / `SaveMemoryAndEndConversationTool` | `update_identity` / `update_soul` / `save_memory_and_end_conversation` |
|
||||
| `proxy_tools.py` | `SetProxyTool` / `GetProxyTool` | `set_proxy` / `get_proxy` |
|
||||
| `misc_tools.py` | `AskUserTool` | `ask_user` |
|
||||
|
||||
**测试**:[tests/test_b_tools.py](../../aibox/tests/test_b_tools.py) 58 项全过
|
||||
|
||||
### 2.3 Wave 1 子系统抽离(`subsystems/` 目录)
|
||||
|
||||
| 文件 | 类 | 行数 |
|
||||
|---|---|---|
|
||||
| `executors.py` | `CommandExecutor` + `SubprocessCommandExecutor` | 279 |
|
||||
| `prompt_manager.py` | `PromptManager`(单例,数字前缀排序加载) | 241 |
|
||||
| `db.py` | `MemoryDatabase` + `MemoDatabase`(SQLite) | 530 |
|
||||
| `skill_manager.py` | `SkillManager` | 186 |
|
||||
|
||||
**测试**:
|
||||
- [tests/test_executors.py](../../aibox/tests/test_executors.py)
|
||||
- [tests/test_prompt_manager.py](../../aibox/tests/test_prompt_manager.py)
|
||||
- [tests/test_db.py](../../aibox/tests/test_db.py)
|
||||
- [tests/test_skill_manager.py](../../aibox/tests/test_skill_manager.py)
|
||||
|
||||
共 70 项全过。
|
||||
|
||||
### 2.4 工具系统基础设施
|
||||
|
||||
- [tools/base.py](../../aibox/tools/base.py)(115 行):`Tool` 抽象基类,定义 `get_name` / `get_description` / `get_parameters` / `execute` / `execute_streaming` / `validate_params` / `format_result` 接口
|
||||
- [tools/registry.py](../../aibox/tools/registry.py)(135 行):`ToolRegistry` 注册器,支持热重载与流式输出分发
|
||||
- [tools/__init__.py](../../aibox/tools/__init__.py) / [subsystems/__init__.py](../../aibox/subsystems/__init__.py):采用"重导出"策略,外部统一用 `from tools import Tool, ReadFileTool`,调用点无需关心子模块位置
|
||||
|
||||
### 2.5 关键 Bug 修复
|
||||
|
||||
- **`subsystems/db.py` 的 `get_recent_memories(limit=None)` SQL datatype mismatch**:`LIMIT ?` 传 `None` 时 SQLite 报错;修复为 `limit` 为 `None` 时移除 `LIMIT` 子句
|
||||
- **`test_basic_tools.py` 路径漂移**(2026-08-14 修复):4 个测试用 `filepath="aibox.py"` 假设 pytest 从 `aibox/` 目录运行;改用 pytest `tmp_path` fixture 创建临时文件,测试不再依赖 cwd
|
||||
|
||||
## 3. 量化指标
|
||||
|
||||
| 指标 | 起点(v2.8) | 当前(v3.0) | 变化 |
|
||||
|---|---|---|---|
|
||||
| `aibox.py` 行数 | 8450 | 3706 | -56% |
|
||||
| 抽离工具类数 | 0 | 17 | +17 |
|
||||
| 抽离子系统数 | 0 | 4 | +4 |
|
||||
| 测试用例数 | 0 | 171 | +171 |
|
||||
| 测试运行时长 | — | 10.79s | — |
|
||||
|
||||
## 4. 测试套件现状
|
||||
|
||||
```
|
||||
aibox/tests/ (171 项,10.79s)
|
||||
├── test_basic_tools.py 43 项 A 批工具(文件/网页/时间/计算器)
|
||||
├── test_b_tools.py 58 项 B 批工具(13 个管理面工具)
|
||||
├── test_executors.py — 命令执行器
|
||||
├── test_prompt_manager.py — 提示词管理器
|
||||
├── test_db.py — 记忆/备忘录数据库
|
||||
└── test_skill_manager.py — 技能管理器
|
||||
```
|
||||
|
||||
运行命令:
|
||||
|
||||
```bash
|
||||
# 从仓库根目录
|
||||
py -m pytest aibox/tests -q
|
||||
```
|
||||
|
||||
## 5. 剩余工作(未抽离)
|
||||
|
||||
### 5.1 aibox.py 剩余类清单(3706 行)
|
||||
|
||||
| 类 | 起始行 | 批次 | 说明 |
|
||||
|---|---|---|---|
|
||||
| `ConversationCompressor` | L70 | Wave 3 | 对话压缩器,依赖 `llm_client` |
|
||||
| `TaskScheduler` | L570 | Wave 2 | 任务调度器,依赖 `MemoDatabase` |
|
||||
| `ConversationHistory` | L986 | Wave 2 | 对话历史管理 |
|
||||
| `ExpressEmotionTool` | L1309 | C 批 | 情绪表达工具,依赖 `run_light_subagent` |
|
||||
| `SubAgentTool` | L1574 | C 批 | 子代理派生工具,依赖 `run_full_subagent` |
|
||||
| `DeepSeekChat` | L1700 | 核心类 | 保留入口(编排逻辑) |
|
||||
| `WebSocketHandler` | L2715 | 核心类 | 保留入口 |
|
||||
| `EnhancedAIManager` | L3377 | 核心类 | 保留入口 |
|
||||
|
||||
### 5.2 建议优先级
|
||||
|
||||
| 优先级 | 批次 | 内容 | 风险 |
|
||||
|---|---|---|---|
|
||||
| 高 | Wave 2 | `TaskScheduler` + `ConversationHistory` 抽离至 `subsystems/` | 低(依赖已抽离的 `MemoDatabase`) |
|
||||
| 中 | Wave 3 | `ConversationCompressor` 抽离至 `subsystems/` | 中(需评估与 `DeepSeekChat` 的循环依赖,参考 lessons learned) |
|
||||
| 中 | C 批 | `ExpressEmotionTool` + `SubAgentTool` 抽离至 `tools/` | 中(依赖 `DeepSeekChat` 的 `run_light_subagent` / `run_full_subagent`,需通过依赖注入或协议接口打破循环) |
|
||||
| 低 | 核心类 | `DeepSeekChat` / `WebSocketHandler` / `EnhancedAIManager` | 保留在 `aibox.py`,作为聚合入口 |
|
||||
|
||||
### 5.3 风险点
|
||||
|
||||
1. **`ConversationCompressor` 与 `DeepSeekChat` 的循环依赖**:压缩器需调用 LLM,而 `DeepSeekChat` 持有压缩器;需通过依赖注入或协议接口打破(参考 project_memory 中 lessons learned)
|
||||
2. **C 批工具依赖 `DeepSeekChat` 内部方法**:`run_light_subagent` / `run_full_subagent` 是 `DeepSeekChat` 的方法,抽离工具时需先抽离这些方法或改为函数式调用
|
||||
3. **测试覆盖盲区**:C 批工具(`ExpressEmotionTool` / `SubAgentTool`)和 Wave 2/3 子系统尚无单元测试,抽离时应同步补测
|
||||
|
||||
## 6. 文档同步状态
|
||||
|
||||
| 文档 | 状态 | 备注 |
|
||||
|---|---|---|
|
||||
| [aibox/README.md](../../aibox/README.md) | ✅ 已同步 v3.0 | 行数、目录结构、开发指南、v3.0 更新日志已更新 |
|
||||
| 本报告 | ✅ 新建 | 归档已完成工作与剩余规划 |
|
||||
| [aibox_orchestrator_worker_plan.md](aibox_orchestrator_worker_plan.md) | 历史 | v2.9 分层方案(已实施) |
|
||||
| [aibox_open_dev_plan.md](aibox_open_dev_plan.md) | 待评估 | 开放版裁剪方案,与本次模块化正交 |
|
||||
| [aibox_file_network_cli_refactor.md](aibox_file_network_cli_refactor.md) | 待评估 | 文件/网络/CLI 重构计划 |
|
||||
|
||||
## 7. 验证记录
|
||||
|
||||
```
|
||||
$ py -m pytest aibox/tests -q
|
||||
........................................................................ [ 42%]
|
||||
........................................................................ [ 84%]
|
||||
........................... [100%]
|
||||
171 passed in 10.79s
|
||||
```
|
||||
|
||||
- 测试基线:2026-08-14
|
||||
- 运行环境:Windows + PowerShell 5
|
||||
- Python:通过 `py` 启动器调用(`python` 不在 PATH)
|
||||
@@ -1,355 +0,0 @@
|
||||
# 认证独立化(CA 模块)开发文档
|
||||
|
||||
**日期**: 2026-08-09(v2 重写)
|
||||
**背景**: gateway v0.8.0 单文件 [main.cpp](../../gateway/src/main.cpp)(2470 行),认证逻辑(约 65%)与反向代理耦合;core/dnfs2 通过 HTTP 回调 gateway 验签,gateway 既是流量入口又是认证权威,形成环形依赖
|
||||
**范围**: `ca`(新模块)、`gateway`(瘦身为纯代理);**WAF / core / dnfs2 / static / 前端零改动**
|
||||
**前版**: 2026-08-08 初版计划书(auth 双栈方案)。本版按已选定拓扑重写:**数据流 `waf → ca → gateway`,gateway 不再验签**。原双栈/后置方案见 §2.2 决策记录
|
||||
|
||||
---
|
||||
|
||||
## 0. 摘要
|
||||
|
||||
新建 CA(Central Auth)模块:C++20 服务,接管 gateway 当前监听的 `127.0.0.1:8443`(WAF upstream 不变),承载**全部认证面**(登录/注册/Token 生命周期/CAPTCHA 代理/限流/WS 代理),并作为**认证前置层**:对业务请求本地验签后注入可信身份头,转发给退居 `127.0.0.1:8445` 的 gateway。gateway 删除全部认证代码,不再验签,仅保留白名单反代。
|
||||
|
||||
核心收益:验签在 CA 本地完成(RSA 密钥同进程,**零新增网络跳**);core/dnfs2 仍调 `8443` 的 `/admin/verify`、`/whoami`,**无感知切换**;前端 URL 不变;现有黑盒 34 断言原样充当验收门。
|
||||
|
||||
---
|
||||
|
||||
## 1. 现状分析
|
||||
|
||||
### 1.1 当前调用链
|
||||
|
||||
```
|
||||
客户端 ──HTTPS──▶ WAF:443 ──HTTP──▶ gateway:8443 ──HTTP──▶ static:9200 / dnfs2:8006 / core:5200
|
||||
│ ├─ 本地验签(RSA 内存密钥)+ Cookie 回退 + Authorization 注入
|
||||
│ ├─ /login /register /logout /refresh_token /whoami /admin/verify
|
||||
│ ├─ /register/send_code /reset/send_code /reset_password
|
||||
│ ├─ GET /captcha/image(pymsg 代理 captcha2:8009)
|
||||
│ └─ WS 代理 /ws/aibox(握手验签 is_admin)
|
||||
└── pymsg: db:8005 / captcha2:8009 / smtp3:8008 / confs:8010
|
||||
|
||||
下游验签回调(HTTP 回环):
|
||||
core:5200 ──HTTP──▶ gateway:8443/admin/verify (HttpTask.cpp,GATEWAY_HOST/GATEWAY_PORT 可配,fail-closed)
|
||||
dnfs2:8006 ──HTTP──▶ gateway:8443/whoami (gateway_url 配置项,fail-closed)
|
||||
```
|
||||
|
||||
### 1.2 gateway 职责盘点与归属
|
||||
|
||||
| 职责域 | 函数/组件 | 归属 |
|
||||
|--------|-----------|------|
|
||||
| 用户认证 | `handleLogin` / `handleRegister` / `handleLogout` / `handleResetPassword` / `handleSendEmailCode` / `handleResetSendCode` | → CA |
|
||||
| Token 管理 | `verify_token` / `generate_jti` / `revoke_jti` / `revoke_all_user_tokens` / `handleRefreshToken` / `handleWhoami` / `handleAdminVerify` / jti 黑名单 / 用户版本号表 / RSA 密钥对 | → CA |
|
||||
| 密码体系 | `PBKDF2Hasher` / `PasswordValidator` / 防枚举 dummy 校验 | → CA |
|
||||
| 验证码/邮箱 | `handleCaptchaImage` / `verify_captcha_remote` / `verify_email_code_remote` | → CA |
|
||||
| 速率限制 | 全部 RateLimiter 实例(登录 5/5min、注册 3/10min、发码 5/10min + 邮箱级 3/h+10/d、验证码 10/min 等) | → CA |
|
||||
| WS 代理 | `verifyWsAdminToken` + `_WS_PROXY_ROUTES` 注册循环 | → CA(验签+桥接一体) |
|
||||
| Cookie | `q4_token` 签发/清除、`extract_token_from_request`(Cookie 回退)、代理时 `Authorization: Bearer` 注入 | → CA |
|
||||
| HTTP 反代 | `handleProxy`(白名单/路径安全/头转发/XFF) | 留 gateway(去验签化) |
|
||||
| 头卫生 | `is_special_header` / `is_safe_target` | 留 gateway |
|
||||
| 基础设施 | `wrapperHttp` / `ConfIni` / `DbC` / `LogC` / 信号处理 / 优雅关闭 | 两侧各持一份(模板见 smtp3) |
|
||||
|
||||
### 1.3 核心问题
|
||||
|
||||
1. **职责耦合**:认证与代理同进程,任一域改动爆炸半径为全文件
|
||||
2. **环形依赖**:gateway 是流量入口又是验签权威,core/dnfs2 须 HTTP 回调入口
|
||||
3. **敏感面混杂**:RSA 私钥、密码哈希校验、jti 黑名单与转发逻辑共用内存空间
|
||||
4. **不可独立扩展**:认证密集与代理密集负载无法分别扩容;认证状态纯内存,绑定单实例
|
||||
|
||||
---
|
||||
|
||||
## 2. 目标架构
|
||||
|
||||
### 2.1 拓扑
|
||||
|
||||
```
|
||||
客户端 ──HTTPS──▶ WAF:443 ──HTTP──▶ CA:8443 ──HTTP──▶ gateway:8445 ──▶ static:9200 / dnfs2:8006 / core:5200
|
||||
│
|
||||
├─ 本地处理:/login /register /logout /refresh_token /whoami /admin/verify
|
||||
│ /register/send_code /reset/send_code /reset_password
|
||||
│ GET /captcha/image
|
||||
├─ WS:/ws/aibox 握手验签(is_admin) → 直接桥接 ws://127.0.0.1:8765
|
||||
├─ 其余路径:本地验签 → 注入 X-Auth-User/Authorization → 转发 gateway:8445
|
||||
└─ pymsg: db:8005 / captcha2:8009 / smtp3:8008 / confs:8010
|
||||
|
||||
端口安排:CA 接管 8443(WAF upstream 配置不变);gateway 迁 8445(仅 loopback,仅 CA 可达)。
|
||||
```
|
||||
|
||||
职责划分:**CA = 认证权威 + 认证前置代理**;**gateway = 纯路由/转发策略**(opened 白名单、路径安全、XFF 透传)。
|
||||
|
||||
### 2.2 决策记录
|
||||
|
||||
| 方案 | 拓扑 | 结论 |
|
||||
|------|------|------|
|
||||
| **选定:CA 前置** | `waf → ca → gateway` | ✅ 验签在 CA 本地(零网络跳);core/dnfs2 调 8443 不变;WAF/前端不变;gateway 无需新增路径路由机制 |
|
||||
| 否决 A:CA 后置双栈 | `waf → gateway → ca`(HTTP+pymsg 双栈) | gateway 需新增路径级路由表;CORS/Set-Cookie/错误映射下沉 CA;两服务改动面都大 |
|
||||
| 否决 B:CA 后置纯 pymsg | `waf → gateway →(pymsg) ca` | gateway 保留 HTTP↔pymsg 翻译层与 Cookie 职责,认证面仍在入口进程内,瘦身不彻底;且每请求多一跳 pymsg 验签或引入吊销延迟的缓存 |
|
||||
| 否决 C:验签缓存 | gateway 缓存验签结果 60s | logout/改密后 Token 在窗口内仍可用,吊销语义弱化(随拓扑选定自然消失) |
|
||||
|
||||
### 2.3 信任模型(重要)
|
||||
|
||||
1. **loopback 信任基线不变**:本架构 loopback 内部服务本就不互鉴权(db:8005 execute 无认证即先例)。gateway:8445 仅绑 `127.0.0.1`,与本基线一致。
|
||||
2. **身份头契约(strip-then-set)**:CA 在边缘剥离客户端伪造的 `X-Auth-User` / `X-Auth-Admin`(Beast fields 大小写不敏感,`erase` 后 `set`),验签通过才注入。模式 A 下 CA 的唯一上游是 WAF(loopback)。
|
||||
3. **gateway 存在性不变量**:显式 `Target-Url` 请求必须携带 `X-Auth-User`,否则 401(非验签,仅防线编程,防本地误直连/SSRF 兜底)。
|
||||
4. **模式 B 警示迁移**:CA 直连公网(无 WAF)时 XFF 可被伪造、限流可绕过——原 gateway README 的 A4 降级提示随限流一并迁入 CA README;根治方案 [libhttp_remote_addr_dev_doc.md](libhttp_remote_addr_dev_doc.md) 中的 `trust_xff`/`rate_limit_key()` 配套改在 CA 落地。
|
||||
|
||||
---
|
||||
|
||||
## 3. CA 模块设计
|
||||
|
||||
### 3.1 目录与构建(沿用 smtp3/captcha2 约定)
|
||||
|
||||
```
|
||||
ca/
|
||||
├── src/main.cpp # 入口 + HTTP 路由 + 认证逻辑 + 前置转发(迁移代码为主)
|
||||
├── inc/
|
||||
│ ├── Conf.hpp # ConfIni(confs 拉取 + env 覆盖 + 完整性校验)
|
||||
│ ├── DbC.hpp # 自 gateway 随迁(已基于 libmsg us/Msg.hpp;client 名 "gateway"→"ca")
|
||||
│ ├── LogC.hpp # 异步日志(smtp3 长连接版)
|
||||
│ ├── ExpBase.hpp # 异常基类
|
||||
│ └── LocalTime.hpp # 时间工具
|
||||
├── Makefile # 链接 us_http / us_msg / us_crypt / ssl / crypto / boost_json / curl
|
||||
├── _set.sh / _run.sh
|
||||
└── README.md
|
||||
```
|
||||
|
||||
- 二进制 `ca.bin`;版本从 `1.0.0` 起;日志服务名 `ca`
|
||||
- 依赖:libhttp(HttpServer/HttpClient/ws_proxy)、libmsg(MsgClient/RpcClient)、libcrypt(MyRSA/Base64)、OpenSSL(PBKDF2/RAND_bytes)、boost::json/program_options/property_tree、libcurl
|
||||
- **不依赖 libsqlite3**:Users 表访问全部经 db:8005(现状即如此)
|
||||
|
||||
### 3.2 配置(`confs/configs/ca.ini` 新增)
|
||||
|
||||
```ini
|
||||
[Server]
|
||||
address = 127.0.0.1
|
||||
port = 8443 # 接管 gateway 现端口,WAF upstream 不变
|
||||
|
||||
[Upstream]
|
||||
gateway = http://127.0.0.1:8445 # 业务请求唯一上游
|
||||
|
||||
[Token] # 自 gateway.ini 原样迁入
|
||||
keep_time = 3600
|
||||
refresh_grace_multiplier = 2
|
||||
max_refresh = 10
|
||||
|
||||
[Admin] # 原样迁入
|
||||
users = 1776177662@qq.com;277665039@qq.com;
|
||||
|
||||
[WsProxy] # 自 gateway.ini 迁入
|
||||
routes = /ws/aibox=ws://127.0.0.1:8765;
|
||||
|
||||
[CORS] # 原样迁入(P2-3:预检反射由 libhttp 处理,白名单暂仍不生效)
|
||||
allowed_origins = https://qsont.me;https://www.qsont.me;http://127.0.0.1:8443;http://localhost:8443
|
||||
|
||||
[Downstream] # 原样迁入
|
||||
db_host = 127.0.0.1
|
||||
db_port = 8005
|
||||
captcha_host = 127.0.0.1
|
||||
captcha_port = 8009
|
||||
smtp_host = 127.0.0.1
|
||||
smtp_port = 8008
|
||||
```
|
||||
|
||||
环境变量覆盖(与 gateway 现机制一致):`LISTEN_ADDR`/`LISTEN_PORT`、`UPSTREAM_GATEWAY`、`DB_*`、`CAPTCHA_*`、`SMTP_*`、`CONFS_HOST`/`CONFS_PORT`。配置加载:confs `get` 拉取(retry_on_refused 阻塞)→ env 覆盖 → 完整性校验。
|
||||
|
||||
### 3.3 HTTP 路由表(对外契约与原 gateway 完全一致)
|
||||
|
||||
| 路由 | 处理 | 说明 |
|
||||
|------|------|------|
|
||||
| POST `/login` | 本地 | 强制远端 CAPTCHA;dummy PBKDF2 防枚举;签发 Token + `q4_token` Cookie |
|
||||
| POST `/register` | 本地 | 密码强度→邮箱格式→CAPTCHA→邮箱验证码(smtp fail-closed)→写 Users |
|
||||
| POST `/register/send_code`、`/reset/send_code` | 本地 | IP+邮箱级限流、强制 CAPTCHA、smtp 60s 冷却 429 透传 |
|
||||
| POST `/reset_password` | 本地 | 防枚举恒成功文案;改密后版本号递增全设备吊销 |
|
||||
| POST `/logout` | 本地 | jti 吊销 + 清 Cookie |
|
||||
| POST `/refresh_token` | 本地 | 宽限期刷新、旧 jti 吊销、max_refresh 上限 |
|
||||
| POST `/whoami` | 本地 | 恒 200,`{login:false}` 不泄露细节;core/dnfs2 依赖此契约 |
|
||||
| POST `/admin/verify` | 本地 | fail-closed 401;core 依赖此契约 |
|
||||
| GET `/captcha/image` | 本地 | pymsg 调 captcha2,PNG + `X-Captcha-Token` + `no-store`,10/min/IP |
|
||||
| WS `[WsProxy]` 路径 | 本地 | 握手验签 `is_admin=true` → 桥接后端(ws/wss),双向透传 |
|
||||
| GET/POST `*`(通配) | **转发** | 见 §3.4 |
|
||||
|
||||
### 3.4 前置转发层(通配路由 → gateway:8445)
|
||||
|
||||
迁移 `handleProxy` 的转发骨架,认证语义前置:
|
||||
|
||||
1. **路径安全**:`is_safe_target()` 校验(`..`/`%2e`/`%2f`/`%5c`/空字节 → 400),CA/gateway 双层各执行一次(防线叠加,gateway 侧保留原校验)。
|
||||
2. **公开 vs 受保护**:无 `Target-Url` 头 → 公开目标(Proxy.default 语义在 gateway),**剥离**客户端 `X-Auth-User`/`X-Auth-Admin` 后直接转发,不注入任何凭证。
|
||||
3. **受保护**:有 `Target-Url` → 提取 Token(`Authorization` 优先,回退 `q4_token` Cookie)→ `verify_token` 本地验签(RSA→过期→jti→版本号)→ 失败按现契约 400/401,**不转发**;成功则注入:
|
||||
- `Authorization: Bearer <token>`(Cookie 来源也统一补齐,dnfs2 等下游零改动)
|
||||
- `X-Auth-User: <email>`、`X-Auth-Admin: true`(仅管理员;供 gateway 不变量与下游未来免回调使用)
|
||||
4. **头卫生**:剥离客户端自带的 `X-Auth-User`/`X-Auth-Admin`(strip-then-set);XFF 透传 WAF 值(缺省 `unknown`);`Host`/`connection`/`Accept-Encoding` 等逐跳头处理沿用现 `is_special_header` 集合 + `x-auth-user`/`x-auth-admin` 新增。
|
||||
5. **转发失败语义**:gateway:8445 不可达 → 502(与现 captcha/smtp fail-closed 风格一致),记 warning 日志。
|
||||
6. **响应**:状态码与响应头原样回传(现 `handleProxy` 行为),`Connection: close`。
|
||||
|
||||
### 3.5 迁移资产清单(gateway `main.cpp` → CA `main.cpp`)
|
||||
|
||||
代码以**整体搬迁**为主,保持实现逐字不变(降低回归面),仅做两处适配:日志服务名 `gateway`→`ca`;DbC client 名同上。
|
||||
|
||||
| 资产 | 备注 |
|
||||
|------|------|
|
||||
| PBKDF2Hasher / PasswordValidator / 常数时间比较 | 含 SHA256 旧哈希迁移注释块(保持注释状态) |
|
||||
| TokenInfo / verify_token / generate_jti / revoke_jti / revoke_all_user_tokens / jti 黑名单(5000 上限清理)/ 用户版本号表 | RSA 密钥对启动内存生成(重启即全吊销,语义不变) |
|
||||
| 8 个 RateLimiter 实例 + 邮箱级双窗口限流 | 限流键仍取 XFF(模式 A 可信;模式 B 警示见 §2.3-4) |
|
||||
| handleLogin ~ handleResetPassword 全部 handler | 错误文案/状态码契约逐字保留(黑盒断言依赖) |
|
||||
| handleWhoami / handleAdminVerify / handleRefreshToken / handleLogout | 同上 |
|
||||
| handleCaptchaImage / verify_captcha_remote / verify_email_code_remote | fail-closed 语义保留 |
|
||||
| verifyWsAdminToken + ws_proxy 注册 | 桥接目标直接为后端(不再经 gateway) |
|
||||
| handleProxy 转发骨架 / is_safe_target / is_special_header / get_cookie_value / extract_token_from_request | 按 §3.4 改造 |
|
||||
| wrapperHttp / 信号处理 / 优雅关闭(8s+3s 超时)/ ExpBase / ConfIni / LocalTime / DbC / LogC | 基础设施随迁 |
|
||||
|
||||
### 3.6 不变量自检(启动时)
|
||||
|
||||
- `[Upstream] gateway` 非空且为 `http://127.0.0.1:*`(防误配到公网)
|
||||
- 模式 B 检测:`Server.address` 非 loopback 时启动日志打印 A4 安全降级警告
|
||||
|
||||
---
|
||||
|
||||
## 4. gateway 瘦身设计(v0.9.0)
|
||||
|
||||
### 4.1 删除清单(预计 -1700 行,2470 → ~750 行)
|
||||
|
||||
- 全部认证 handler 与 Token/PBKDF2/RSA/jti/版本号/限流/CAPTCHA/邮箱验证码代码
|
||||
- `q4_token` Cookie 读写、`extract_token_from_request`、代理内验签与 Authorization 注入
|
||||
- `[Token]`/`[Admin]`/`[WsProxy]`/`[Database]` 配置读取;下游 captcha/smtp 配置(db 保留供 LogC)
|
||||
- `/login` 等 9 条认证路由 + `/captcha/image` 路由 + WS 代理注册循环
|
||||
|
||||
### 4.2 保留清单
|
||||
|
||||
- `handleProxy`(去验签化)+ `is_safe_target` + `is_special_header` + `wrapperHttp`
|
||||
- 存在性不变量:显式 `Target-Url` 且无 `X-Auth-User` → 401(§2.3-3)
|
||||
- 公开目标规则:无 `Target-Url` → `Proxy.default`,不校验身份头
|
||||
- XFF 处理:透传 CA 值(缺省 `unknown`)
|
||||
- 配置加载、信号处理、优雅关闭、LogC
|
||||
|
||||
### 4.3 精简后 `confs/configs/gateway.ini`
|
||||
|
||||
```ini
|
||||
[Server]
|
||||
address = 127.0.0.1
|
||||
port = 8445 # 迁出 8443;仅 CA 可达
|
||||
tls = off
|
||||
|
||||
[Proxy]
|
||||
default = http://127.0.0.1:9200
|
||||
closed = http://127.0.0.1:8005;http://127.0.0.1:8007;
|
||||
opened = http://127.0.0.1:9200;http://127.0.0.1:8006;http://127.0.0.1:5200;
|
||||
|
||||
[Downstream] # 仅供 LogC 日志落库
|
||||
db_host = 127.0.0.1
|
||||
db_port = 8005
|
||||
```
|
||||
|
||||
### 4.4 启动自检
|
||||
|
||||
- `Server.address` 必须为 loopback(`127.0.0.1`/`::1`):gateway 不再验签后,绑定公网等于裸奔,**非 loopback 直接拒绝启动**并打印致命错误。
|
||||
|
||||
---
|
||||
|
||||
## 5. 接口契约
|
||||
|
||||
### 5.1 对外 HTTP(不变量)
|
||||
|
||||
`https://<host>/login`、`/register`、`/whoami` 等全部 URL、请求/响应 JSON 结构、状态码(400/401/403/429/500/502)、`q4_token` Cookie 属性、`X-Captcha-Token` 头——**逐字不变**。前端、WAF、core、dnfs2 零改动的依据即此。
|
||||
|
||||
### 5.2 内部头契约(CA ↔ gateway ↔ 下游)
|
||||
|
||||
| 头 | 方向 | 语义 |
|
||||
|----|------|------|
|
||||
| `X-Auth-User` | CA→gateway→下游 | 已验签身份;CA strip-then-set;gateway 对显式 Target-Url 要求其存在 |
|
||||
| `X-Auth-Admin` | CA→gateway→下游 | `true` 表示管理员;供下游免回调豁免(可选消费) |
|
||||
| `Authorization` | CA→gateway→下游 | Cookie 来源统一补齐为 Bearer;下游解析不变 |
|
||||
| `X-Forwarded-For` | WAF→CA→gateway→下游 | 逐层透传,缺省 `unknown` |
|
||||
| `Target-Url` | 客户端→CA→gateway | 路由策略头,语义不变 |
|
||||
|
||||
### 5.3 Token 格式(不变)
|
||||
|
||||
`Base64(payload)_RSA-Sign(payload)`;payload 字段(user/jti/ver/iat/created_at/keep_time/is_admin/refresh_count/max_refresh)不变。切换瞬间旧 Token 全部失效(CA 新密钥),用户重新登录——一次性事件,cutover 通告中说明。
|
||||
|
||||
### 5.4 下游零改动验证点
|
||||
|
||||
| 消费方 | 现状调用 | 切换后 |
|
||||
|--------|----------|--------|
|
||||
| core `verify_admin_token` | `http://127.0.0.1:8443/admin/verify` | 同 URL 由 CA 服务,契约不变 |
|
||||
| dnfs2 Token 验签 | `http://127.0.0.1:8443/whoami` | 同上 |
|
||||
| static `/api/check_user` | 直连 db 查 Users(账号旁路) | 不受影响;遗留项另行收口 |
|
||||
|
||||
---
|
||||
|
||||
## 6. 上线步骤(cutover 与回滚)
|
||||
|
||||
| 步骤 | 内容 |
|
||||
|------|------|
|
||||
| 1 | CA 开发完成,`LISTEN_PORT=18443` 并行实例联调(不影响现网 8443) |
|
||||
| 2 | gateway 瘦身版(v0.9.0)构建,`gateway.ini` 改 8445 |
|
||||
| 3 | cutover(秒级窗口):停 gateway:8443 → 启 CA:8443 → 启 gateway:8445 → 黑盒全量验证 |
|
||||
| 4 | `core/init.txt` 增加 `ca/_run.sh`(gateway 之前);README/INDEX 文档同步 |
|
||||
| 5 | 观察 24h:CA/gateway 日志、db Logs 表 4xx/5xx 基线对比 |
|
||||
|
||||
**回滚**:停 CA → 旧 gateway 二进制 + 旧配置回 8443 → 重启。Token 全部重签(用户重新登录),无数据迁移。
|
||||
|
||||
---
|
||||
|
||||
## 7. 测试计划
|
||||
|
||||
### 7.1 验收门(零改动复用)
|
||||
|
||||
- `tests/blackbox/test_gateway.py` 34 断言**原样**对 8443 执行:认证路径打到 CA、代理路径经 CA→gateway,外部契约不变则全绿(当前基线 32 passed / 2 env-skipped)
|
||||
- 修复期两个红测试(A3 缺字段 400)已在 v0.8.0 转绿,迁移后必须保持
|
||||
|
||||
### 7.2 新增探针(test_ca.py,黑盒补充)
|
||||
|
||||
| 探针 | 预期 |
|
||||
|------|------|
|
||||
| 客户端伪造 `X-Auth-User: admin@x` + 显式 Target-Url(无 Token) | CA 剥离后 401,不到下游 |
|
||||
| 伪造 `X-Auth-User` + 有效非管理员 Token 访问 WS 路由 | 403(以验签结果为准,不以头为准) |
|
||||
| 直连 gateway:8445,显式 Target-Url、无 `X-Auth-User` | 401(存在性不变量) |
|
||||
| 经 CA 全链路(8443)带 Token 访问 dnfs2 | 下游收到 `X-Auth-User` 与 `Authorization` |
|
||||
| CA 存活、gateway:8445 停止 | 业务路径 502(不静默放行),认证路径正常 |
|
||||
| gateway 存活、CA 停止 | WAF 502(入口不可达,符合预期 fail-closed) |
|
||||
| 公开目标(无 Target-Url)带伪造身份头 | 下游收不到 `X-Auth-User`/`Authorization` |
|
||||
|
||||
### 7.3 集成回归
|
||||
|
||||
- core HTTP 管理接口(`/admin/verify` 链路)实测;dnfs2 笔记 CRUD(`/whoami` 链路 + AI 审核)实测
|
||||
- WS 代理端到端:管理员 Token 升级 `/ws/aibox` 双向帧;非管理员 403
|
||||
- 限流全谱在 CA 复测(随机 XFF 技巧沿用);cutover/回滚演练各一次
|
||||
|
||||
### 7.4 测试资产演进(后续)
|
||||
|
||||
- `test_gateway.py` 按进程边界拆分为 `test_ca.py`(A/B/C/E 段)+ 精简 `test_gateway.py`(D 段代理矩阵 + 8445 直连探针),README/INDEX 同步
|
||||
- 认证 happy path 仍依赖 captcha2/smtp 测试钩子(既有遗留,CA 落地后可仅在 CA 层加 `[Security] captcha_required` 开关,单点解锁)
|
||||
|
||||
---
|
||||
|
||||
## 8. 风险与缓解
|
||||
|
||||
| 风险 | 等级 | 缓解 |
|
||||
|------|------|------|
|
||||
| CA 不可用 → 全站入口中断 | 高 | CA 职责内聚后可用性等同于现 gateway;core 守护 + 黑盒监控;回滚方案秒级 |
|
||||
| 迁移遗漏导致认证行为漂移 | 中 | 代码逐字搬迁 + 34 断言验收门 + §7.2 探针矩阵 |
|
||||
| 身份头伪造/泄露 | 中 | strip-then-set + gateway 存在性不变量 + 公开目标不注入 + 直连探针回归 |
|
||||
| gateway:8445 被公网暴露 | 高 | 启动强制 loopback 自检(§4.4);部署审计 checklist 增加端口扫描项 |
|
||||
| cutover 窗口 Token 全失效 | 低 | 预期内(RSA 换钥),通告用户重新登录 |
|
||||
| 双重转发延迟 | 低 | 本机 HTTP 转发 <1ms;CA→gateway 短连接可后续优化为 keep-alive |
|
||||
| 模式 B 限流绕过(A4 迁移) | 中 | CA README 降级警示;根治按 libhttp_remote_addr_dev_doc(配套改落 CA) |
|
||||
|
||||
---
|
||||
|
||||
## 9. 工作量估算
|
||||
|
||||
| 部分 | 规模 | 估时 |
|
||||
|------|------|------|
|
||||
| CA 脚手架(目录/Makefile/_set.sh/_run.sh/LogC 等) | 模板复制 + 适配 | 0.5 天 |
|
||||
| 认证资产搬迁 + 前置转发层改造 | ~1700 行搬迁 + ~150 行改造 | 1.5 天 |
|
||||
| gateway 瘦身 + 启动自检 | -1700 行 + ~20 行 | 0.5 天 |
|
||||
| 配置/init.txt/文档(README×2、INDEX、confs) | — | 0.5 天 |
|
||||
| 黑盒验收 + §7.2 探针 + 集成回归 + cutover 演练 | — | 1 天 |
|
||||
| **合计** | | **4 天**(含缓冲) |
|
||||
|
||||
---
|
||||
|
||||
## 10. 与既有文档的关系
|
||||
|
||||
- 本文档**替代** 2026-08-08 初版计划书中的双栈拓扑与阶段划分(初版的职责盘点/迁移清单仍然有效,已并入 §1.2/§3.5)
|
||||
- [libhttp_remote_addr_dev_doc.md](libhttp_remote_addr_dev_doc.md)(A4 根治):实施时 `trust_xff`/`rate_limit_key()` 配套由 gateway 改为落 CA;`is_special_header` 两侧同步加 `x-remote-addr`
|
||||
- [gateway_SECURITY_AUDIT.md](../security/gateway_SECURITY_AUDIT.md):A2(白名单 403)断言、P2-3(CORS 白名单)随转发层/配置迁入 CA;P1-7(CSRF Token)在 CA 层评估
|
||||
- gateway README / 项目 README 的架构图、端口表、请求链路在上线步骤第 4 步同步更新
|
||||
@@ -312,13 +312,13 @@ systemctl start Nidus # 或 bash core/_run.sh
|
||||
| **core ConfsTask** | C++ 内嵌配置任务(557 行) |
|
||||
| **gateway 精简** | `src/main.cpp` 大幅精简(-1980 行),认证逻辑下沉至 ca |
|
||||
|
||||
### 8.2 认证独立化(auth_service_plan)
|
||||
### 8.2 认证独立化(ca 模块)
|
||||
|
||||
`waf → ca → gateway` 前置拓扑:ca 统一鉴权并前置转发,gateway 退化为纯转发(仅 loopback),下游 / WAF / 前端零改动。对应开发文档见 `docs/reports/auth_service_plan.md`。
|
||||
`waf → ca → gateway` 前置拓扑:ca 统一鉴权并前置转发,gateway 退化为纯转发(仅 loopback),下游 / WAF / 前端零改动。
|
||||
|
||||
### 8.3 关键 dev docs 索引(`docs/reports/`)
|
||||
|
||||
core_service_registry_dev_doc(注册中心)、auth_service_plan(CA 模块)、aibox_open_dev_plan / aibox_file_network_cli_refactor(aibox 安全沙箱与 CLI 化)、libcrypt_CBC_deprecation_dev_doc(AES-GCM 迁移)、libhttp_remote_addr_dev_doc(X-Remote-Addr 透传)、captcha_cpp_rewrite / smtp_cpp_rewrite / confs2_cpp_rewrite(C++ 重写方案)。
|
||||
core_service_registry_dev_doc(注册中心)、aibox_open_dev_plan(aibox 安全沙箱与 CLI 化)、libcrypt_CBC_deprecation_dev_doc(AES-GCM 迁移)、libhttp_remote_addr_dev_doc(X-Remote-Addr 透传)、smtp_cpp_rewrite(C++ 重写方案)。
|
||||
|
||||
### 8.4 安全评级演进
|
||||
|
||||
|
||||
@@ -1,406 +0,0 @@
|
||||
# bot4cj 模块 Python 迁移开发方案
|
||||
|
||||
> 文档类型:功能开发计划
|
||||
> 模块:bot4cj → pybot(Python 迁移)
|
||||
> 作者:自动生成
|
||||
> 日期:2026-08-09
|
||||
|
||||
---
|
||||
|
||||
## 1. 概述
|
||||
|
||||
### 1.1 背景与动机
|
||||
|
||||
`bot4cj` 是基于仓颉(Cangjie)编程语言和 Linux I2C 接口的机器人硬件脚本解析器,通过 `.mt` 矩阵脚本描述机器人硬件构成与动作序列,驱动 PCA9685 舵机控制板、MPU6050 传感器等真实硬件(宿主机为树莓派 aarch64)。
|
||||
|
||||
将其迁移到 Python 的动机:
|
||||
|
||||
| 收益 | 说明 |
|
||||
|------|------|
|
||||
| 工具链门槛消除 | 仓颉编译器(cjc 1.1.3)+ cjpm + 4 个仓颉扩展包(ctype4cj/pthread4cj/spi4cj/lua4cj)→ 树莓派系统自带 Python3,零编译 |
|
||||
| 与 actpred 生态打通 | actpred(动作预测,PyTorch)已有 Python 版 `.mt` 解析器(`mt_parser.py`),迁移后"大脑(actpred)↔ 身体(bot4cj)"共享同一语言与数据结构,DRL 闭环开发成本大幅降低 |
|
||||
| Lua 扩展层可消除 | 仓颉版引入 Lua 是为了运行时可扩展硬件驱动;Python 本身即脚本语言,硬件驱动直接写成 Python 插件模块,砍掉 lua4cj + liblua5.4 + C 函数注册整层 |
|
||||
| C/C++ 共享库大部分可消除 | libi2c4cj.so 的功能可用 `fcntl.ioctl` + `os.read/write` 纯 Python 实现;矩阵/FFT/CNN 由 numpy/PyTorch(actpred 已有)承担 |
|
||||
| 审计遗留一并修复 | [bot4cj 代码审计](../reviews/bot4cj_CODE_AUDIT_20260808.md) 15 项发现(3 P0 + 5 P1 + 7 P2)全部待修复,迁移时逐项闭环 |
|
||||
| 与项目主流技术栈对齐 | micro_servers 业务服务以 Python/C++20 为主,仓颉是唯一的孤立技术栈,维护人手受限 |
|
||||
|
||||
**性能评估**:bot4cj 的执行节奏由硬件决定(舵机 50Hz PWM、I2C 100kHz~400kHz、动作行间 20~500ms 延时),解析与调度不是瓶颈,Python 的性能完全够用;且 Python 插件替代 Lua 桥接后,单次硬件调用的框架开销反而下降约两个数量级(详见附录 A 的量化分析)。这与 db2/captcha2 等"Python→C++ 求性能"的方向相反,是"仓颉→Python 求生态与可维护性"。
|
||||
|
||||
### 1.2 设计目标
|
||||
|
||||
1. **`.mt` 脚本零改动**:现有 `scripts/*.mt` 全部原样可用,块语法(`@Alias`/`@I2C`/`@PCA9685`/`@MPU6050_BIND`/`@Speaker`/`@Servo`/`@Action`)、`-` 占位符语义、`#` 注释、`;` 行结束符完全兼容。
|
||||
2. **CLI 兼容**:保持 `cat scripts/std.mt | ./bot` 的 stdin 管道用法与 `--help`/`--test` 分支。
|
||||
3. **零第三方依赖的核心**:解析器 + 硬件层仅用标准库(`fcntl`/`os`/`struct`/`threading`/`json`/`argparse`),树莓派上无需 venv 即可运行;numpy 仅测试/工具链可选。
|
||||
4. **硬件插件化(Python 替代 Lua)**:未知 `@Block` 路由到 `modules/<Name>.py` 插件,保留 build/call/destroy 三段生命周期协议。
|
||||
5. **审计 15 项发现全部修复**:见 §5 修复映射表。
|
||||
6. **可测试**:pytest 单元测试覆盖解析器、占位符语义、Servo 校准、插件加载;硬件层可用 Mock I2C 在无硬件的 x86 上跑通全流程。
|
||||
|
||||
### 1.3 不做什么
|
||||
|
||||
- **不迁移 math/NN/CNN/FFT/CUDA 层**(`src/math/*.cj`、`src/cpp/cnn/`、`src/cpp/fft*`、`src/cuda/`):这部分是为 DRL 预留的自研神经网络库,功能已被 actpred(PyTorch)完全覆盖且质量更高。DRL 闭环未来直接调 actpred,不重复造轮子。相应地 `--test nn/cnn/cnn1` 分支不迁移。
|
||||
- **不迁移 `.bot` 旧版脚本支持**:`scripts/*.bot` 是历史遗留,仓颉版本身也未解析它。
|
||||
- **不改 `.mt` 语法**:`@Action_Auto` 等语法演进(旧分析 D11)留待后续版本。
|
||||
- **不引入 smbus2/RPi.GPIO 等三方库**:I2C 用 `fcntl.ioctl(fd, I2C_SLAVE=0x0703, addr)` 直接实现,与仓颉版 `I2C.cj`/`I2C.c` 行为逐字节对齐。
|
||||
- **不删除 bot4cj 目录**:遵循项目"原版与重写版并存 → 验证通过后切换"的渐进策略(同 db→db2、captcha→captcha2、smtp→smtp3)。
|
||||
|
||||
---
|
||||
|
||||
## 2. 现状分析
|
||||
|
||||
### 2.1 仓颉版架构分层
|
||||
|
||||
```
|
||||
.mt 矩阵脚本(硬件定义 + 动作序列)
|
||||
↓
|
||||
MatrixParser(P.cj,FSM 逐字符解析)→ MatrixObject 链表
|
||||
↓
|
||||
RobotStruct(r.cj)createHardVars():按 @块名分发到硬件变量工厂
|
||||
MatrixVars_Hardware 基类 + 三闭包:rise/things/free
|
||||
↓
|
||||
硬件实现 hd/*.cj:I2C(20 槽锁数组)、Servo、PCA9685、MPU6050、Speaker
|
||||
+ Lua 扩展 modules/*.lua(未知 @Block 的通用驱动,经 C 函数注册调 I2C/SPI)
|
||||
↓
|
||||
系统库 .so:libi2c4cj(I2C glibc 封装)、libsnd4cj(音频)、libmt4cj/libfft4cj/libcnn4cj(数学)
|
||||
```
|
||||
|
||||
### 2.2 迁移范围源码清单
|
||||
|
||||
| 仓颉源文件 | 行数 | 职责 | 迁移去向 |
|
||||
|-----------|------|------|---------|
|
||||
| `src/main.cj` | 178 | CLI 入口、stdin 管道、--test 分支 | `pybot/main.py` |
|
||||
| `src/P.cj` | 264 | MatrixParser(FSM)+ MatrixObject + MatrixSaver + 解析异常 | `pybot/parser.py` |
|
||||
| `src/r.cj` | 1072 | RobotStruct + 硬件变量基类 + 6 个工厂 + Lua 状态机桥接 | `pybot/robot.py` + `pybot/hardvars.py` |
|
||||
| `src/hd/I2C.cj` + `src/cpp/I2C.c` | 340+ | I2C glibc 封装 + 每 fd 互斥锁 | `pybot/hd/i2c.py` |
|
||||
| `src/hd/Servo.cj` | 31 | max/min/offset 校准 + 负 max 角度翻转 | `pybot/hd/servo.py` |
|
||||
| `src/hd/MPU6050.cj` | 244 | 6 轴传感器 + 后台采样线程 + 卡尔曼 | `pybot/hd/mpu6050.py` |
|
||||
| `src/hd/snd.cj` + `src/cpp/snd.cpp` | 23+ | 扬声器(仓颉侧仅 sleep 占位) | `pybot/hd/speaker.py` |
|
||||
| `src/hd/PCA9685.cj` | — | 已标记"由 Lua 取代" | 不迁移(走插件) |
|
||||
| `modules/*.lua`(PCA9685/ADS1115/ADS1256/SPI/TEST) | — | Lua 硬件驱动 | `pybot/modules/*.py` 插件 |
|
||||
| `src/math/*.cj`、`src/cpp/cnn|fft|Matrix|kalman`、`src/cuda/` | ~2000 | 自研 NN/数学库 | **不迁移**(actpred 承担);卡尔曼滤波例外,MPU6050 需要,~40 行纯 Python |
|
||||
|
||||
### 2.3 关键语义(必须逐项对齐)
|
||||
|
||||
1. **占位符 `-`**:递归沿用上一行同列值;首行为 `-` 时,`parseMt` 解析为 `0`(数值列),`getMt` 返回字面 `-`。
|
||||
actpred `mt_parser.py` 曾与此不同(首行 `-` 用 `@Servo` max 列填充),2026-08 已统一为本语义(填 0):运行时 Servo 有 max/min 钳位保护,且 max 列是钳位上限而非初始位姿(反装舵机 max 为负值,会填出指令空间 0~180 之外的非法值)。
|
||||
2. **@Action 列名 `dev.param`**:点号拆分为硬件名 + 参数名;参数名经 `@Alias` 表转换(`A`→`angle`、`F`→`frequency`)后传入硬件。
|
||||
3. **每行执行**:同一行按硬件聚合参数字典,逐硬件调用一次 `play()`;行间时延由 `spk.delay` 参数承担。
|
||||
4. **硬件生命周期**:`active()` 按创建顺序 build,`deactive()` 倒序 destroy,且必须在 `try/finally` 中保证执行。
|
||||
5. **Servo 角度算法**:`a = angle + offset; a = clamp(a, |min|, |max|); if max < 0: a = 180 - a`。
|
||||
6. **PCA9685 脉宽映射**:50Hz、0°→0.5ms、180°→2.5ms、`offval = width * 4096 / 20`、寄存器 `0x06 + 4*channel` 写 4 字节。
|
||||
7. **MPU6050**:±500°/s、±8g、10ms 采样周期、卡尔曼滤波(Q=20, R=200)、roll/pitch 由加速度计反正切求得。
|
||||
8. **I2C 互斥**:同一 fd 上"设地址 + 读/写寄存器"必须是原子序列(多设备共享总线)。
|
||||
|
||||
### 2.4 审计遗留(迁移时修复的输入清单)
|
||||
|
||||
来自 [bot4cj_CODE_AUDIT_20260808.md](../reviews/bot4cj_CODE_AUDIT_20260808.md):3 项 P0(S01 列数校验恒假 / S02 调试 print 污染 stdout / S03 扬声器无实声)、5 项 P1(S04 转义不一致 / S05 MPU6050 地址重复 / S06 有符号读寄存器 / S07 Lua 路径硬编码 / S08 文档滞后)、7 项 P2(无测试、注释代码、锁扫描 O(n)、无编译告警、消费端报错不可读、线程退出、GC 析构不确定)。
|
||||
|
||||
---
|
||||
|
||||
## 3. 目标架构
|
||||
|
||||
### 3.1 目录结构
|
||||
|
||||
```
|
||||
pybot/
|
||||
├── pybot/
|
||||
│ ├── __init__.py
|
||||
│ ├── main.py # CLI 入口(argparse:--help / --test / -i / -o 预留)
|
||||
│ ├── parser.py # MatrixParser + MatrixObject + MatrixSaver + 解析异常
|
||||
│ ├── robot.py # RobotStruct(createHardVars / active / run / deactive / save)
|
||||
│ ├── hardvars.py # HardwareVar 基类 + Alias/I2C/Servo/Speaker/MPU6050_BIND 工厂
|
||||
│ ├── plugin.py # Python 插件加载器(替代 Lua 状态机桥接)
|
||||
│ └── hd/
|
||||
│ ├── __init__.py
|
||||
│ ├── i2c.py # fcntl.ioctl I2C + 每 fd 锁(dict + threading.Lock)
|
||||
│ ├── servo.py # 舵机校准与翻转
|
||||
│ ├── mpu6050.py # 传感器 + threading.Thread 采样 + kalman
|
||||
│ ├── kalman.py # 一维卡尔曼滤波(~40 行)
|
||||
│ └── speaker.py # 正弦波发声(可选 alsaaudio,默认 aplay 子进程 / 静默降级)
|
||||
├── modules/ # 硬件驱动插件(Python 替代 Lua)
|
||||
│ ├── PCA9685.py
|
||||
│ ├── ADS1115.py
|
||||
│ ├── ADS1256.py
|
||||
│ └── SPI.py
|
||||
├── scripts -> ../bot4cj/scripts # 软链或复制,.mt 脚本共享
|
||||
├── tests/
|
||||
│ ├── test_parser.py
|
||||
│ ├── test_hardvars.py
|
||||
│ ├── test_servo.py
|
||||
│ └── mock_i2c.py # Mock I2C 总线(记录寄存器写序列,供断言)
|
||||
├── _run.sh # cat x.mt | bash _run.sh(exec python3 -m pybot)
|
||||
├── _set.sh # 无编译;可选建 venv 装 pytest
|
||||
├── requirements.txt # 核心为空;[dev] pytest
|
||||
└── README.md
|
||||
```
|
||||
|
||||
### 3.2 数据流(与仓颉版一致)
|
||||
|
||||
```
|
||||
stdin → MatrixParser.parse()(逐块产出 MatrixObject)
|
||||
→ RobotStruct.create_hardvars()(按块名分发工厂;未知块 → plugin.load("modules/<Name>.py"))
|
||||
→ robot.active()(按序 build)
|
||||
→ robot.run()(逐行 @Action:聚合参数 → 逐硬件 play(dict))
|
||||
→ robot.deactive()(倒序 destroy,try/finally 保证)
|
||||
```
|
||||
|
||||
### 3.3 关键映射决策
|
||||
|
||||
| 仓颉版机制 | Python 版方案 | 理由 |
|
||||
|-----------|--------------|------|
|
||||
| `JsonValue` 统一通信协议 | 原生 `dict`(进程内直接传递) | 无跨语言边界后无需序列化;插件协议文档化 dict 键约定 |
|
||||
| 三闭包 rise/things/free | `HardwareVar` 基类三方法 `build()/play(params)/destroy()`,子类覆写 | Python 惯用法;闭包捕获改为实例属性 |
|
||||
| Lua 状态机 + C 函数注册(i2c_open 等 11 个) | `importlib` 动态加载 `modules/<Name>.py`,插件直接 `from pybot.hd import i2c` | 消除整个 FFI 层;插件与核心同语言 |
|
||||
| 20 槽 I2C 锁数组(O(n) 线性扫描) | `dict[fd → threading.Lock]` + 全局注册锁 | 修复审计 Q03;容量不再受限 |
|
||||
| pthread_create/pthread_cancel(MPU6050 采样) | `threading.Thread(daemon=True)` + `threading.Event` 停止标志 | 修复审计 Q06/D3,安全退出 |
|
||||
| 仓颉 GC 终结器不确定 | 上下文管理器 `with RobotStruct(...) as robot:` + `atexit` 兜底 | 修复审计 Q07/D1,确定性释放 fd |
|
||||
| cjpm build + Makefile .so | 无构建步骤 | 树莓派上 `git pull` 即部署 |
|
||||
|
||||
### 3.4 插件协议(modules/*.py,替代 Lua build/call/destroy)
|
||||
|
||||
```python
|
||||
# modules/PCA9685.py —— 与 PCA9685.lua 逐语义对齐
|
||||
def build(ctx: dict) -> dict:
|
||||
"""ctx: 元数据(fd/address/... 来自 @PCA9685 块行 + 引用硬件的 play() 展开)
|
||||
初始化:重启 → SLEEP → 设 prescale(50Hz) → 唤醒"""
|
||||
def call(ctx: dict) -> dict:
|
||||
"""ctx 额外含实时参数(channel/angle);写 4 个 PWM 寄存器"""
|
||||
def destroy(ctx: dict) -> dict:
|
||||
"""复位 PCA9685"""
|
||||
```
|
||||
|
||||
- 插件查找顺序(修复审计 S07):`$PYBOT_MODULES_PATH`(冒号分隔多路径)→ 包内置 `modules/` → 当前工作目录 `./modules/`。
|
||||
- 加载失败抛出含插件名与搜索路径的 `PluginNotFound` 异常。
|
||||
|
||||
---
|
||||
|
||||
## 4. 详细设计
|
||||
|
||||
### 4.1 parser.py
|
||||
|
||||
- **MatrixObject**:`name: str` + `rows: list[list[str]]`;提供 `__getitem__((row, col))`、`index_of_col(row, key)`、`save(writer)`。越界访问抛 `MatrixIndexError(块名, 行, 列)`(修复审计 Q05 的报错可读性)。
|
||||
- **MatrixParser**:保持 FSM 逐字符解析(非正则),状态与仓颉版一致(nl/cm/nm 三标志),原因:`@` 块嵌套注释、行内注释、任意空白分隔的兼容性已被现有脚本验证。
|
||||
- 修复 S01:`submit_line` 用 `len(grp) != len(rows[-1])` 校验列数,不一致抛 `DoesNotSync(行号, 期望列数, 实际列数)`。
|
||||
- 修复 S02:删除逐字符 debug 打印;提供 `--verbose` 时输出到 **stderr**(stdout 保留给未来管道通信)。
|
||||
- 修复 S04:Python 无 Rune/String 混淆问题,换行统一 `"\n"` 字符比较;补充单测覆盖"漏写分号跨行"场景。
|
||||
- **占位符语义**(§2.3-1):`get_mt(obj, row, col)` 递归向上;`parse_mt(obj, row, col, type)` 首行 `-` 返回 `type("0")`。
|
||||
- **MatrixSaver**:保持与仓颉版相同的输出格式(`@Name` / 空格分隔 / `;` 结尾 / `@` 收尾),保证 `--test matrix_o/matrix_2` 的解析↔序列化往返可用。
|
||||
|
||||
### 4.2 hd/i2c.py
|
||||
|
||||
```python
|
||||
I2C_SLAVE = 0x0703
|
||||
class I2CBus:
|
||||
def __init__(self, path): self.fd = os.open(path, os.O_RDWR); _register_lock(self.fd)
|
||||
def set_addr(self, addr): fcntl.ioctl(self.fd, I2C_SLAVE, addr)
|
||||
def write_reg(self, reg, val): os.write(self.fd, bytes([reg, val]))
|
||||
def read_reg(self, reg) -> int: # 返回 0..255 无符号(修复 S06)
|
||||
os.write(self.fd, bytes([reg])); return os.read(self.fd, 1)[0]
|
||||
def transaction(self): # with bus.transaction(): 原子序列(设地址+读写)
|
||||
```
|
||||
|
||||
- 锁注册表:`_LOCKS: dict[int, threading.Lock]` + `_MAIN_LOCK`;`transaction()` 上下文管理器持锁,替代仓颉版 `I2cLock_lock/unlock` 手动配对。
|
||||
- `read_reg` 天然返回无符号 int(`bytes[0]` ∈ 0..255),16 位拼接 `(hh << 8) | ll` 后用 `struct.unpack('>h')` 转有符号——修复 S06 的符号扩展错误且保留传感器数据的正确符号。
|
||||
|
||||
### 4.3 hd/servo.py
|
||||
|
||||
与 `Servo.cj` 逐行对齐(§2.3-5),构造参数 `max/min/offset/set_angle_callback`。补单测:正装(max=90)、反装(max=-180 翻转)、offset 漂移、越界钳位。
|
||||
|
||||
### 4.4 hd/mpu6050.py
|
||||
|
||||
- 寄存器常量与 `MPU6050.cj` 一致;**修复 S05**:`ADDR_0 = 0x68`、`ADDR_1 = 0x69`。
|
||||
- 采样线程:`threading.Thread` + `self._stop = threading.Event()`,循环 `while not self._stop.wait(0.01)`(10ms 周期且可即时打断)——修复 Q06。
|
||||
- 读传感器整体包在 `bus.transaction()` 中(设地址 + 14 字节连读),比仓颉版逐寄存器 12 次 I2C 往返更高效且原子。
|
||||
- 卡尔曼滤波参数保持 Q=20、R=200;roll/pitch 公式不变。
|
||||
- 数据出口:`latest()` 返回带 `threading.Lock` 保护的最近一帧滤波值(仓颉版 mutex 字段从未真正使用,此处补全)。
|
||||
|
||||
### 4.5 hd/speaker.py(决策 S03)
|
||||
|
||||
仓颉版 `play()` 仅 sleep(实声 C++ 函数被注释)。Python 版决策:**实现真实发声,带静默降级**。
|
||||
|
||||
- 首选:标准库 `wave` + `subprocess.run(["aplay", "-q", tmp.wav])`——树莓派自带 alsa-utils,零 pip 依赖;正弦波生成用 `math.sin`(每次 ≤ 数秒 @48kHz,纯 Python 可接受),保留 5ms 淡入淡出。
|
||||
- `aplay` 不存在或播放失败:log 一次警告后仅 sleep(delay),行为与仓颉版一致(延时职责不丢失)。
|
||||
- `delay` 语义不变:`play()` 阻塞 `delay` 毫秒(发声与延时并行时取延时为准)。
|
||||
|
||||
### 4.6 robot.py / hardvars.py
|
||||
|
||||
- `RobotStruct` 保持仓颉版结构:`C: list[MatrixObject]` + `M: dict[str, HardwareVar]`(Python dict 有序,天然满足"倒序 destroy"需求)。
|
||||
- `create_hardvars()` 分发表:`{"I2C": I2CVar, "Alias": AliasVar, "Servo": ServoVar, "Speaker": SpeakerVar, "MPU6050_BIND": MPU6050Var}`;`Action` 跳过;未知块 → `plugin.PluginVar`。
|
||||
- 重名键检查(仓颉版 FIXME r.cj:911):同名硬件变量重复定义时,先 `destroy()` 旧值再覆盖,并输出 stderr 警告。
|
||||
- `run()` 逐行执行 `@Action`;列名合法性(必须含一个 `.`)在**执行前**整表校验一次,报错含列号与列名(配合 Q05 修复)。
|
||||
- 上下文管理器:`__enter__`→`active()`,`__exit__`→`deactive()`;`main.py` 用 `with` 替代仓颉版 try/finally。
|
||||
|
||||
### 4.7 main.py
|
||||
|
||||
| 仓颉版 CLI | Python 版 | 说明 |
|
||||
|-----------|-----------|------|
|
||||
| (无参)stdin 管道 | 同 | `cat std.mt | ./bot` 不变;另支持 `pybot path/to/x.mt` 位置参数(增强,非破坏) |
|
||||
| `--help` | argparse 自动生成 | 文案更新 |
|
||||
| `--test nn/cnn/cnn1` | **移除**(见 §1.3) | NN 由 actpred 承担 |
|
||||
| `--test matrix_o/matrix_2` | `--test roundtrip <in.mt> <out.mt>` | 解析↔序列化往返测试,路径改为参数(仓颉版硬编码 `ss/test.mt`) |
|
||||
| `-i` / `-o`(DRL 预留) | 保留参数位,标注 reserved | 未来对接 actpred |
|
||||
| (新增)`--dry-run` | 解析 + 构建硬件变量但不开 I2C,Mock 总线打印动作序列 | 无硬件调试 .mt 脚本 |
|
||||
| (新增)`--verbose` | FSM 与硬件调用日志到 stderr | 替代 S02 删掉的调试打印 |
|
||||
|
||||
### 4.8 日志规范
|
||||
|
||||
- 所有诊断输出走 `logging`(stderr),stdout 保留纯净——为未来 actpred 管道集成预留(审计 §六 的阻塞项 S02 根治)。
|
||||
|
||||
---
|
||||
|
||||
## 5. 审计发现修复映射表
|
||||
|
||||
| 审计 ID | 级别 | 仓颉版问题 | Python 版修复方式 | 落点 |
|
||||
|---------|------|-----------|------------------|------|
|
||||
| S01 | P0 | submitLine 列数校验恒假 | `len(grp) != len(rows[-1])` 显式校验 + DoesNotSync 带期望/实际列数 | parser.py |
|
||||
| S02 | P0 | parse() 逐字符 print 污染 stdout | 移除;--verbose 走 stderr logging | parser.py / main.py |
|
||||
| S03 | P0 | 扬声器仅 sleep 无实声 | wave + aplay 实现发声,缺 aplay 静默降级为 sleep | hd/speaker.py |
|
||||
| S04 | P1 | r"\n" 与 r'\n' 转义不一致 | Python 无此坑;补"漏分号跨行"单测防回归 | tests/test_parser.py |
|
||||
| S05 | P1 | MPU6050 ADDR_1 应为 0x69 | `ADDR_1 = 0x69` | hd/mpu6050.py |
|
||||
| S06 | P1 | readI2cReg 返回有符号 Int8 | read_reg 返回 0..255;16 位拼接后 struct 显式转有符号 | hd/i2c.py |
|
||||
| S07 | P1 | Lua 路径硬编码 ./modules/ | `$PYBOT_MODULES_PATH` → 包内置 → cwd 三级查找 | plugin.py |
|
||||
| S08 | P1 | README 已知问题表滞后 | pybot README 重写;bot4cj README 标注维护状态 | README.md |
|
||||
| Q01 | P2 | 无单元测试 | pytest:解析器/占位符/Servo/插件/Mock I2C 全流程 | tests/ |
|
||||
| Q02 | P2 | 大量注释旧代码 | 新代码库不携带 | — |
|
||||
| Q03 | P2 | I2C 锁 O(n) 线性扫描 | dict[fd → Lock] | hd/i2c.py |
|
||||
| Q04 | P2 | Makefile 告警关闭 | 无 C 编译;CI 跑 `python -m pyflakes`(可选) | _set.sh |
|
||||
| Q05 | P2 | 列数不匹配报错不可读 | MatrixIndexError 带块名/行/列 + run() 前整表校验 | parser.py / robot.py |
|
||||
| Q06 | P2 | pthread_cancel 非安全点退出 | threading.Event 优雅停止 | hd/mpu6050.py |
|
||||
| Q07 | P2 | GC 析构不确定,fd 泄漏 | with 上下文管理器 + atexit 兜底 | robot.py / main.py |
|
||||
|
||||
---
|
||||
|
||||
## 6. 测试计划
|
||||
|
||||
### 6.1 单元测试(无硬件,x86 可跑)
|
||||
|
||||
| 用例组 | 覆盖 |
|
||||
|--------|------|
|
||||
| test_parser | 全部 `scripts/*.mt` 解析成功;列数不一致抛 DoesNotSync;漏分号跨行;注释/空行/CRLF;roundtrip(解析→save→再解析等价) |
|
||||
| test_placeholder | `-` 递归沿用;首行 `-` 数值列=0;actpred mt_parser 已统一为同一语义(2026-08),两侧测试互为语义锚点 |
|
||||
| test_servo | 正装/反装翻转/offset/钳位边界 |
|
||||
| test_hardvars | @Alias 转换;@Action 列名拆分;同行参数聚合;重名覆盖警告 |
|
||||
| test_plugin | 插件三级路径查找;PluginNotFound 报错内容;PCA9685.py 用 Mock I2C 断言寄存器写序列(prescale=121、reg=0x06+4*ch、offval 字节序) |
|
||||
| test_mpu6050 | Mock 总线注入 14 字节样本 → 有符号拼接正确(高字节 ≥ 0x80 场景,回归 S06);Event 停止 1s 内退出 |
|
||||
|
||||
### 6.2 实机验证(树莓派 + Dashun 蜘蛛机器人)
|
||||
|
||||
1. `cat scripts/std.mt | bash _run.sh`:站立姿态与仓颉版目测一致(10 舵机角度逐一对照 --verbose 日志)。
|
||||
2. `jump.mt` / `lie-down.mt` / `tilt-*.mt` 逐个回放,动作时序无可感知差异。
|
||||
3. 双 MPU6050(0x68 + 0x69)同总线并行采样(回归 S05),roll/pitch 数值与仓颉版单设备输出同量级。
|
||||
4. 扬声器 `spk.F` 发声可闻(回归 S03),无爆音(淡入淡出生效)。
|
||||
5. Ctrl-C 中断:舵机停止、I2C fd 全部释放(`ls /proc/<pid>/fd` 验证)、MPU6050 线程退出。
|
||||
|
||||
### 6.3 交叉验证
|
||||
|
||||
- `--dry-run` 模式下用 Mock 总线记录全部 I2C 写序列,与仓颉版(打补丁记录)同脚本输出逐字节 diff——作为迁移等价性的最强证据(可参考 captcha2 迁移的像素级验证经验)。
|
||||
|
||||
---
|
||||
|
||||
## 7. 里程碑
|
||||
|
||||
| 阶段 | 内容 | 产出 |
|
||||
|------|------|------|
|
||||
| M1 解析器 | parser.py + 全量单测 + roundtrip | 全部 scripts/*.mt 解析通过 |
|
||||
| M2 硬件抽象 | hardvars.py + robot.py + plugin.py + Mock I2C | --dry-run 全流程可跑(x86) |
|
||||
| M3 硬件落地 | i2c.py / servo.py / mpu6050.py / speaker.py + modules/*.py | 树莓派实机 std.mt 站立 |
|
||||
| M4 等价性验证 | §6.2 实机回放 + §6.3 写序列 diff | 验证报告 |
|
||||
| M5 切换 | README/INDEX 更新、bot4cj 标注 legacy、aibox/BotOps 若有引用改指 pybot | 并存期结束评估 |
|
||||
|
||||
---
|
||||
|
||||
## 8. 风险与对策
|
||||
|
||||
| 风险 | 影响 | 对策 |
|
||||
|------|------|------|
|
||||
| Python GIL + 解释器抖动影响动作时序 | 舵机动作节奏轻微不均 | 动作行间延时 ≥ 20ms,抖动 <1ms 不可感知;MPU6050 采样用 Event.wait 定时,不忙等;量化依据见附录 A |
|
||||
| aplay 子进程首次启动延迟 | 首个音符滞后 | build() 时预热(播放 0 长度静音);或后续可选接 alsaaudio |
|
||||
| 仓颉版行为的未记录细节 | 等价性缺口 | §6.3 写序列逐字节 diff 兜底 |
|
||||
| .mt 脚本方言分裂(pybot vs actpred 解析器) | 语义漂移 | 占位符语义已统一(2026-08,首行 `-` 均填 0);长期可评估 actpred 复用 pybot parser 的只读模式 |
|
||||
| 树莓派 Python 版本参差 | 语法兼容 | 目标 Python ≥ 3.9(树莓派 OS Bullseye 起自带 3.9) |
|
||||
|
||||
---
|
||||
|
||||
## 9. 与 actpred 的集成(已实现:进程内闭环)
|
||||
|
||||
原规划的"大脑-身体"闭环已在 `run_loop()` 中落地,采用进程内调用替代 FIFO 序列化:
|
||||
|
||||
```
|
||||
camera 帧 → actpred.predict(F帧, H历史) → 动作向量(18) → pybot 下发执行
|
||||
↑ ↓
|
||||
history deque ←────────── 回写实际动作 ←────────────────────┘
|
||||
```
|
||||
|
||||
**用法**:`python3 -m pybot --loop CKPT --hz 10 scripts/std.mt`
|
||||
|
||||
- checkpoint 后缀 `.hef` → Hailo NPU 后端,否则 → PyTorch 后端
|
||||
- `@Action` 表头列数 ≤ `action_dim`(默认 18),列索引即动作向量索引
|
||||
- 闭环中剔除 `spk.delay` 阻塞(发声但不延时),节拍由 `hz` 主循环掌控
|
||||
- 冷启动:F 帧重复同帧 + H 步零动作历史
|
||||
- `--loop-steps N` 限步数调试;`-i/-o` FIFO 保留用于未来跨进程隔离
|
||||
|
||||
**与录制的衔接**:`--record DIR` 产出 actpred 格式 episode → actpred 训练 → 推理 checkpoint → `--loop` 闭环执行,形成完整数据飞轮。
|
||||
|
||||
---
|
||||
|
||||
## 附录 A:性能评估依据
|
||||
|
||||
迁移评审中最常见的两个疑虑是"Python 做驱动会不会慢"与"Python 替代 Lua 会不会引入开销"。本附录给出量化分析,结论:**两者均不构成风险,后者反而是性能收益**。
|
||||
|
||||
### A.1 Python 做硬件驱动的时间尺度分析
|
||||
|
||||
bot4cj 架构中所有微秒级时序都已下沉到专用硬件芯片(PCA9685 自主生成 PWM 波形),主机侧只发送设定值,该路径的瓶颈是 I2C 总线物理速率,与宿主语言无关:
|
||||
|
||||
| 环节 | 耗时量级 | 由谁承担 |
|
||||
|------|---------|---------|
|
||||
| 舵机 PWM 波形(50Hz、0.5~2.5ms 脉宽) | µs 级精度 | PCA9685 硬件芯片,主机不参与 |
|
||||
| I2C 写一个寄存器(100kHz 标准模式) | ~200–300 µs/次 | 总线物理速率,语言无关 |
|
||||
| Python 一次 `os.write`/`ioctl` 调用开销 | ~2–5 µs | 解释器 + 系统调用 |
|
||||
| 设一个舵机角度(4 次寄存器写) | ~1 ms | 总线 ≈ 99%,Python ≈ 1% |
|
||||
| 舵机机械转动 60° | ~100–200 ms | 物理惯性 |
|
||||
| `.mt` 动作行间延时 | 20–500 ms | 脚本显式指定 |
|
||||
|
||||
Python 解释开销比总线时间低约两个数量级,比机械动作低四五个数量级。逐项检查潜在慢点:
|
||||
|
||||
- **MPU6050 10ms 采样循环**:14 字节连读 ≈ 1.5–2ms 总线时间 + 卡尔曼滤波 6 次浮点(纯 Python 数 µs),周期内富余充足;且连读比仓颉版逐寄存器 12 次往返(每次重发从机地址)更快。
|
||||
- **正弦波生成**:48kHz 下每秒音频约 4.8 万次 `math.sin`,纯 Python 10–30ms 算完,发声前一次性生成,可接受。
|
||||
- **GIL/GC**:仅主线程 + 1 个采样线程,GIL 无竞争;GC 停顿亚毫秒级,低于 `sleep()` 本身的 ms 级精度。
|
||||
|
||||
Python 不适用的边界(bot4cj 均不涉及):主机 GPIO bit-bang 生成 PWM/WS2812 单线协议、>1kHz 闭环控制、硬实时保证(Linux 非 RTOS,仓颉版同样做不到)。若未来 DRL 需要高频控制环,正确做法是下沉到协处理器(如 RP2040),属架构问题而非语言问题。
|
||||
|
||||
### A.2 Python 插件替代 Lua 桥接的开销对比
|
||||
|
||||
单论解释器速度 Lua 5.4 快于 CPython,但现有 Lua 路径的开销主要在跨语言桥接而非驱动代码本身。仓颉版每次硬件调用([r.cj MtVars_Lua.things](../../bot4cj/src/r.cj#L311-L331))的实际成本:
|
||||
|
||||
| 步骤 | 实现 | 开销量级 |
|
||||
|------|------|---------|
|
||||
| 元数据深拷贝 | `JsonValue.fromStr(metadata.toJsonString())`——JSON 序列化+反序列化一整轮 | ~10–50 µs |
|
||||
| 参数推入 Lua 表 | 每键 `mallocCString` + `Lua_pushstring/setfield` + `free`,跨仓颉→C→Lua 三层 FFI | ~几 µs/键 |
|
||||
| `Lua_pcallk` + 返回表遍历 | `Lua_next` 循环 + 每键值 `Lua_tostring` 转仓颉 String | ~几 µs |
|
||||
| 强制调试打印 | 仓颉侧 2 次 `println` + Lua 侧 `print_table` 全表打印 | ~几十 µs 至 ms 级(终端 I/O) |
|
||||
| Lua 内 4 次 `i2c_write_reg` | 每次跨 Lua→C→仓颉 @C 三层 | 桥接 ~1 µs + 总线 ~250 µs |
|
||||
|
||||
合计框架开销 ~0.1–1 ms/次,与 I2C 总线时间同量级。Python 插件路径的对应成本:
|
||||
|
||||
| 步骤 | 实现 | 开销量级 |
|
||||
|------|------|---------|
|
||||
| 上下文构造 | `{**metadata, **params}` dict 合并(同进程传引用,无序列化) | ~0.5 µs |
|
||||
| 插件调用 | 同解释器直接函数调用 | ~0.1 µs |
|
||||
| I2C 写 | 插件直接 `i2c.write_reg()`,一层 `os.write` 到底 | 系统调用 ~2 µs + 总线 ~250 µs |
|
||||
| 诊断输出 | 默认关闭,`--verbose` 时走 stderr logging(即审计 S02 修复) | 0 |
|
||||
|
||||
框架开销 ~1–2 µs/次,比 Lua 桥接低约两个数量级。插件内部计算(脉宽换算、移位掩码)在两种解释器上的差距为百纳秒级,相对 1ms 总线时间可忽略。
|
||||
|
||||
**结论**:Lua 层在仓颉版存在的唯一理由是"编译型宿主做不了运行时扩展",为此付出 JSON 往返 + 五层 FFI(仓颉→C→Lua→C→仓颉→glibc)的税;宿主换成 Python 后调用链缩短为 Python→glibc 一层,此项迁移是删除一整层间接性,而非语言互换。
|
||||
|
||||
---
|
||||
|
||||
## 参考文档
|
||||
|
||||
- [bot4cj 代码审计报告](../reviews/bot4cj_CODE_AUDIT_20260808.md)(15 项发现)
|
||||
- [bot4cj 深度分析报告](bot4cj_analysis_20260525_00.md)(D1-D29)
|
||||
- [bot4cj README.robot.md](../../bot4cj/README.robot.md)
|
||||
- [actpred mt_parser.py](../../actpred/mt_parser.py)(Python 解析器先例)
|
||||
|
||||
---
|
||||
|
||||
*本文档为开发计划,实施前需结合实际硬件环境评估。*
|
||||
@@ -1,750 +0,0 @@
|
||||
# captcha 模块 C++20 重写 · 开发方案
|
||||
|
||||
> 文档类型:模块重写开发计划
|
||||
> 模块:captcha(验证码服务)
|
||||
> 日期:2026-08-09
|
||||
> 参照:dnfs2 / db2 / fsys 的 Python → C++20 重写路径
|
||||
|
||||
---
|
||||
|
||||
## 1. 概述
|
||||
|
||||
### 1.1 背景与动机
|
||||
|
||||
captcha 模块是 micro_servers 中唯一仍未 C++ 化的纯 Python 内部基础服务(smtp 同为 Python 但依赖外部 SMTP 会话,重写收益有限)。当前 Python 版核心只有 299 行,依赖仅 Pillow + pymsg,属于**小而精的理想重写目标**。
|
||||
|
||||
| 维度 | Python 版现状 | 重写收益 |
|
||||
|------|-------------|----------|
|
||||
| 依赖 | Python venv + Pillow + pymsg 本地包 | 消除 Python 运行时依赖,编译为单二进制 |
|
||||
| 启动 | `_run.sh` 激活 venv 后 `exec python3` | 与 db2/dnfs2/fsys 一致的 `./captcha2.bin` 直接启动 |
|
||||
| 性能 | PIL 图片生成是 CPU 密集型,单线程 GIL 受限 | 原生 C++ 无 GIL,`workers>0` 时真实并发 |
|
||||
| 审计遗留 | P2 `random.choices` 非安全随机、NEW-2 无速率限制 | 一并用 `std::mt19937` + `random_device` 修复 |
|
||||
| 一致性 | 项目已有 3 个成功的 C++ 重写案例(db2/dnfs2/fsys) | 与 gateway/db2/dnfs2/fsys 统一技术栈 |
|
||||
|
||||
### 1.2 设计目标
|
||||
|
||||
1. **协议完全兼容**:pymsg 三条路由(`image` / `verify` / `health`)请求/响应格式与 Python 版字节级一致,gateway 零改动
|
||||
2. **图片效果一致**:160×50 PNG,相同干扰策略(旋转字符 + 干扰线 + 噪点),视觉效果等同,机器识别难度不降低
|
||||
3. **安全对齐**:修复审计报告中的 `random` → 安全随机、添加生成速率限制
|
||||
4. **零新增重型依赖**:仅引入 FreeType(libfreetype-dev),已经是 Linux 标配;PNG 编码使用单头文件 stb_image_write.h(公共领域)
|
||||
5. **基础设施复用**:libmsg(RpcServer/AsyncQueue)、boost::json、Conf.hpp 模式,所有已在项目内的组件
|
||||
|
||||
### 1.3 与 gateway 内置 CAPTCHA 的关系
|
||||
|
||||
gateway 内部保留了一套旧版本地 CAPTCHA(`generate_captcha_text()` + `verify_captcha()` + `_CAPTCHA_STORE`),与 captcha 服务通过 pymsg 调用的 `verify_captcha_remote()` **并存**。本次重写仅涉及独立的 captcha 服务进程,不影响 gateway 内部逻辑。
|
||||
|
||||
---
|
||||
|
||||
## 2. 现状分析
|
||||
|
||||
### 2.1 源码结构
|
||||
|
||||
```
|
||||
captcha/
|
||||
├── captcha_service.py # 299 行,主程序
|
||||
├── log_client.py # 159 行,异步日志客户端(项目通用 SDK)
|
||||
├── _set.sh # venv 创建 + pip install
|
||||
├── _run.sh # venv 激活 + exec python3
|
||||
└── requirements.txt # ../pymsg + Pillow
|
||||
```
|
||||
|
||||
### 2.2 功能清单
|
||||
|
||||
| 功能 | 代码位置 | 行数 | C++ 等效方案 |
|
||||
|------|---------|------|-------------|
|
||||
| 配置加载(confs → env) | L40–134 | 95 行 | 复用 gateway `Conf.hpp` 的 confs 拉取 + env 覆盖模式 |
|
||||
| `generate_text()` | L145–146 | 2 行 | `std::mt19937` + 字符集数组(审计 NEW-1 修复:用 `random_device` 种子) |
|
||||
| `generate_image()` | L149–182 | 34 行 | FreeType 渲染 + 像素旋转合成 + stb_image_write PNG 编码(约 180 行 C++) |
|
||||
| `CaptchaServer.__init__` | L192–198 | 7 行 | `RpcServer` 构造 + `add_route` 注册三个路由 |
|
||||
| `_handle_image` | L215–241 | 27 行 | 逻辑完全相同:生成 text → 惰性清理 → 容量检查 → 存 store → 调 generate_image → base64 编码返回 |
|
||||
| `_handle_verify` | L243–269 | 27 行 | 逻辑完全相同:pop token → 不存在返回 invalid → XOR 恒定时间比较 |
|
||||
| `_handle_health` | L271–272 | 2 行 | 返回 `{"status":"ok","service":"captcha"}` |
|
||||
| `main` / argparse | L279–299 | 21 行 | boost::program_options(与 gateway 一致) |
|
||||
| log_client | 整个文件 | 159 行 | 基于 libmsg `AsyncQueue`,参照 gateway `LogC.hpp`(约 60 行) |
|
||||
|
||||
### 2.3 `generate_image()` 详细分析
|
||||
|
||||
当前 Python 版图片生成流程(这是 C++ 重写的核心难点):
|
||||
|
||||
```
|
||||
1. Image.new('RGB', (160, 50), random_bg_color)
|
||||
→ 创建 RGB 画布,随机背景色
|
||||
|
||||
2. for each char in text (4 chars):
|
||||
a. random_color → 随机前景色
|
||||
b. random_y (2~12) → 垂直偏移
|
||||
c. random_x (15 + i*32 ± 3) → 水平分布在 160px 宽画布上
|
||||
d. ImageFont.truetype(font_path, random_size 20~38) → 可变字号
|
||||
e. 创建 64×64 透明子画布 → draw_char.text((20,20), char)
|
||||
f. char.rotate(random_angle ±30°) → 旋转字符
|
||||
g. img.paste(rotated_char, (x-32, y-32), mask=rotated_char)
|
||||
→ Alpha 合成到主画布
|
||||
|
||||
3. 4 条随机干扰线(Bresenham)
|
||||
|
||||
4. 50 个随机噪点(单像素)
|
||||
|
||||
5. img.save(buf, format='PNG') → PNG 编码
|
||||
```
|
||||
|
||||
每个步骤的 C++ 等价实现:
|
||||
|
||||
| 步骤 | Python (Pillow) | C++ (FreeType + 手写) |
|
||||
|------|----------------|----------------------|
|
||||
| 1. 创建画布 | `Image.new('RGB', 160, 50)` | `std::vector<uint8_t> pixels(160*50*3)` |
|
||||
| 2a. 随机颜色 | `random.randint(0,256)` | `std::uniform_int_distribution<int>(0, 255)` |
|
||||
| 2d. 字体加载 | `ImageFont.truetype(path, size)` | `FT_New_Face` + `FT_Set_Char_Size` |
|
||||
| 2e. 字符渲染 | `draw_char.text((20,20), ch)` | `FT_Load_Char(face, ch, FT_LOAD_RENDER)` → 灰度位图 |
|
||||
| 2f. 旋转 | `char.rotate(angle, expand=True)` | 手写最近邻插值旋转(~40 行) |
|
||||
| 2g. Alpha 合成 | `img.paste(char, pos, mask=char)` | 像素级 alpha blend(灰度值 → 透明度) |
|
||||
| 3. 干扰线 | `draw.line([...])` | Bresenham 画线算法 |
|
||||
| 4. 噪点 | `draw.point((x, y))` | 直接写像素 |
|
||||
| 5. PNG 编码 | `img.save(buf, format='PNG')` | `stbi_write_png_to_func` |
|
||||
|
||||
### 2.4 依赖替换矩阵
|
||||
|
||||
| Python 依赖 | C++ 替代 | 引入方式 |
|
||||
|------------|---------|----------|
|
||||
| `Pillow` (PIL) | FreeType 2 + stb_image_write.h | `apt install libfreetype-dev` + 单头文件放入 `inc/` |
|
||||
| `pymsg` | libmsg (`RpcServer` / `AsyncQueue`) | 系统已安装 (`libus_msg.so`) |
|
||||
| `log_client.py` | 独立 `LogC.hpp`(基于 `AsyncQueue`) | 参照 gateway `LogC.hpp` 简化版 |
|
||||
| `configparser` | boost::property_tree(`Conf.hpp` 模式) | 系统已安装 (`libboost-all-dev`) |
|
||||
| `argparse` | boost::program_options | 系统已安装 |
|
||||
| `secrets` / `random` | `<random>` (`mt19937` + `random_device`) | C++11 标准库 |
|
||||
| `base64` | libcrypt `Base64` 或手写(~20 行) | 系统已安装 (`libus_crypt.so`) 或内联 |
|
||||
| `threading.Lock` | `std::mutex` + `std::lock_guard` | C++11 标准库 |
|
||||
| `json` | `boost::json` | 系统已安装 |
|
||||
|
||||
---
|
||||
|
||||
## 3. 整体设计
|
||||
|
||||
### 3.1 模块结构
|
||||
|
||||
```
|
||||
captcha2/
|
||||
├── src/
|
||||
│ └── main.cpp # ~420 行,主程序
|
||||
│ ├── 配置加载 # confs captcha.ini → env override(~80 行)
|
||||
│ ├── CaptchaStore # 线程安全 Token 存储(~40 行)
|
||||
│ ├── generate_text() # 安全随机文本生成(~15 行)
|
||||
│ ├── generate_image() # FreeType + 像素操作 + PNG 编码(~180 行)
|
||||
│ ├── base64_encode() # 手写或复用 libcrypt(~20 行)
|
||||
│ └── main() # RpcServer 路由注册 + 启动(~85 行)
|
||||
├── inc/
|
||||
│ ├── stb_image_write.h # 单头文件 PNG 编码器(公共领域,~1800 行)
|
||||
│ └── LogC.hpp # 简化版异步日志客户端(~60 行)
|
||||
├── Makefile # 构建(~30 行)
|
||||
├── _set.sh # make + install(~10 行)
|
||||
├── _run.sh # 启动脚本(~5 行)
|
||||
└── README.md
|
||||
```
|
||||
|
||||
### 3.2 与现有 C++ 基础库的关系
|
||||
|
||||
```
|
||||
captcha2.bin
|
||||
│
|
||||
├── libmsg (RpcServer) ← pymsg 协议 C++ 实现,路由分发 + 响应包装
|
||||
│ └── libus_msg.so (系统已安装)
|
||||
│
|
||||
├── libmsg (AsyncQueue) ← 异步批量日志推送
|
||||
│ └── libus_msg.so
|
||||
│
|
||||
├── boost::json ← JSON 解析/序列化
|
||||
│ └── libboost_json (系统已安装)
|
||||
│
|
||||
├── boost::program_options ← 命令行参数解析
|
||||
│ └── libboost_program_options(系统已安装)
|
||||
│
|
||||
├── FreeType 2 ← 字体渲染
|
||||
│ └── libfreetype.so (需新增,apt install libfreetype-dev)
|
||||
│
|
||||
└── stb_image_write.h ← PNG 编码(header-only,直接放入 inc/)
|
||||
```
|
||||
|
||||
### 3.3 配置体系
|
||||
|
||||
与 gateway/db2/dnfs2/fsys 一致的**两层加载**模式:
|
||||
|
||||
1. **confs 统一配置** → 启动时 pymsg `get` 路由拉取 `captcha.ini`(`retry_on_refused=true` 阻塞等待 confs 就绪)
|
||||
2. **环境变量覆盖**(最高优先级,Docker 兼容)
|
||||
|
||||
新增可选配置项(审计 NEW-2 修复):
|
||||
|
||||
```ini
|
||||
[captcha]
|
||||
len = 4 # 验证码字符数(不变)
|
||||
ttl = 300 # Token 有效期,秒(不变)
|
||||
font_path = /usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf # 字体路径(不变)
|
||||
max_store = 10000 # 最大 Token 存储数(Python 版已有)
|
||||
rate_limit_per_min = 60 # [新增] 每分钟最多生成次数,0=不限制,默认 60
|
||||
```
|
||||
|
||||
配置白名单合并(复用 Python 版的 `_merge_ini` 策略,仅覆盖已知 section/key):
|
||||
|
||||
```cpp
|
||||
struct CaptchaConfig {
|
||||
// [server]
|
||||
std::string host = "127.0.0.1";
|
||||
int port = 8009;
|
||||
int workers = 4; // RpcServer 工作线程数
|
||||
|
||||
// [captcha]
|
||||
int len = 4;
|
||||
int ttl = 300;
|
||||
std::string font_path = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf";
|
||||
int max_store = 10000;
|
||||
int rate_limit_per_min = 60; // 0 = 不限制
|
||||
};
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 4. 详细设计
|
||||
|
||||
### 4.1 CaptchaStore — 线程安全 Token 存储
|
||||
|
||||
```cpp
|
||||
class CaptchaStore {
|
||||
public:
|
||||
struct Entry {
|
||||
std::string answer; // 小写答案
|
||||
std::chrono::steady_clock::time_point created;
|
||||
};
|
||||
|
||||
// 生成新 Token,返回 token 字符串;容量满时抛出 std::runtime_error
|
||||
std::string generate(const std::string& answer, int ttl_seconds, int max_size);
|
||||
|
||||
// 校验并消费 Token(一次性);返回 true 表示答案正确
|
||||
// 无效/过期 Token 返回 false
|
||||
bool verify(const std::string& token, const std::string& answer, int ttl_seconds);
|
||||
|
||||
// 当前存储数量
|
||||
size_t size() const;
|
||||
|
||||
private:
|
||||
void expire_locked(int ttl_seconds); // 惰性清理过期项
|
||||
|
||||
mutable std::mutex mtx_;
|
||||
std::unordered_map<std::string, Entry> store_;
|
||||
};
|
||||
```
|
||||
|
||||
**与 Python 版的关键差异:**
|
||||
- `generate()` 和 `verify()` 都在锁内先清理过期项
|
||||
- Token 生成使用 256-bit 安全随机(`std::random_device` + `std::mt19937` + hex 编码),替代 Python 版 `secrets.token_hex(16)` 的 128-bit → 升级为 256-bit
|
||||
- 容量检查同 Python 版:清理过期后仍超限则拒绝生成
|
||||
|
||||
### 4.2 generate_text() — 安全随机验证码文本
|
||||
|
||||
```cpp
|
||||
std::string generate_text(int len = 4) {
|
||||
// 剔除易混淆字符:l I 1 O 0 o(与 Python 版完全一致)
|
||||
static const char kChars[] =
|
||||
"ABCDEFGHJKLMNPQRSTUVWXYZabcdefghjkmnpqrstuvwxyz23456789";
|
||||
static constexpr size_t kNumChars = sizeof(kChars) - 1;
|
||||
|
||||
// thread_local 保证多线程安全(RpcServer workers>0 时并发回调)
|
||||
thread_local std::mt19937 rng(std::random_device{}());
|
||||
std::uniform_int_distribution<size_t> dist(0, kNumChars - 1);
|
||||
|
||||
std::string text;
|
||||
text.reserve(len);
|
||||
for (int i = 0; i < len; ++i)
|
||||
text.push_back(kChars[dist(rng)]);
|
||||
return text;
|
||||
}
|
||||
```
|
||||
|
||||
**审计修复**:Python 版使用 `random.choices`(确定性伪随机),C++ 版用 `std::random_device` 作为种子源,消除可预测性风险。
|
||||
|
||||
### 4.3 generate_image() — 核心图片生成(FreeType + 像素操作)
|
||||
|
||||
这是重写的**技术核心**,约 180 行 C++。以下为伪代码级设计:
|
||||
|
||||
```cpp
|
||||
// 返回值:PNG 字节流
|
||||
std::vector<uint8_t> generate_image(const std::string& text,
|
||||
const std::string& font_path) {
|
||||
constexpr int kWidth = 160, kHeight = 50, kChannels = 3;
|
||||
std::vector<uint8_t> pixels(kWidth * kHeight * kChannels);
|
||||
|
||||
// 1. 随机背景色
|
||||
fill_random_color(pixels, kWidth, kHeight);
|
||||
|
||||
// 2. 初始化 FreeType
|
||||
FT_Library ft;
|
||||
FT_Init_FreeType(&ft);
|
||||
FT_Face face;
|
||||
// 字体缺失时回退逻辑:先尝试指定路径 → 再尝试默认 DejaVu → 失败则报错退出
|
||||
if (FT_New_Face(ft, font_path.c_str(), 0, &face) != 0)
|
||||
FT_New_Face(ft, "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 0, &face);
|
||||
|
||||
// 3. 逐字符渲染
|
||||
thread_local std::mt19937 rng(std::random_device{}());
|
||||
std::uniform_int_distribution<int> angle_dist(-30, 30);
|
||||
std::uniform_int_distribution<int> size_dist(20, 38);
|
||||
std::uniform_int_distribution<int> y_dist(2, 12);
|
||||
std::uniform_int_distribution<int> x_jitter(-3, 3);
|
||||
|
||||
for (size_t i = 0; i < text.size(); ++i) {
|
||||
// 随机颜色、字号、旋转角度、位置
|
||||
auto color = random_color(rng);
|
||||
int font_size = size_dist(rng);
|
||||
FT_Set_Char_Size(face, 0, font_size * 64, 96, 96);
|
||||
|
||||
int angle = angle_dist(rng);
|
||||
int y = y_dist(rng);
|
||||
int x = 15 + static_cast<int>(i) * 32 + x_jitter(rng);
|
||||
|
||||
// FT_Load_Char 渲染为 8-bit 灰度位图
|
||||
FT_Load_Char(face, text[i], FT_LOAD_RENDER);
|
||||
|
||||
// 旋转灰度位图(最近邻插值,expand 模式)
|
||||
auto rotated = rotate_bitmap(face->glyph->bitmap, angle);
|
||||
|
||||
// Alpha 合成到主画布
|
||||
composite_char(pixels, kWidth, kHeight, rotated, x, y, color);
|
||||
}
|
||||
|
||||
FT_Done_Face(face);
|
||||
FT_Done_FreeType(ft);
|
||||
|
||||
// 4. 4 条随机干扰线
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
draw_line(pixels, kWidth, kHeight,
|
||||
random_point(rng), random_point(rng),
|
||||
random_color(rng));
|
||||
}
|
||||
|
||||
// 5. 50 个随机噪点
|
||||
for (int i = 0; i < 50; ++i) {
|
||||
set_pixel(pixels, kWidth, kHeight,
|
||||
random_point(rng), random_color(rng));
|
||||
}
|
||||
|
||||
// 6. PNG 编码
|
||||
std::vector<uint8_t> png_buf;
|
||||
stbi_write_png_to_func(
|
||||
[](void* ctx, void* data, int size) {
|
||||
auto* buf = static_cast<std::vector<uint8_t>*>(ctx);
|
||||
buf->insert(buf->end(), (uint8_t*)data, (uint8_t*)data + size);
|
||||
},
|
||||
&png_buf, kWidth, kHeight, kChannels,
|
||||
pixels.data(), kWidth * kChannels);
|
||||
|
||||
return png_buf;
|
||||
}
|
||||
```
|
||||
|
||||
**关键辅助函数**:
|
||||
|
||||
| 函数 | 说明 | 预估行数 |
|
||||
|------|------|----------|
|
||||
| `fill_random_color(pixels, w, h)` | 填充随机背景色 | 10 |
|
||||
| `random_color(rng)` | 返回 `(R, G, B)` 结构体 | 5 |
|
||||
| `rotate_bitmap(bitmap, angle)` | 最近邻插值旋转灰度位图,expand 模式 | 40 |
|
||||
| `composite_char(pixels, w, h, rotated, x, y, color)` | 将旋转后字符位图以指定颜色 Alpha 合成到主画布;灰度值作为透明度 | 25 |
|
||||
| `draw_line(pixels, w, h, p1, p2, color)` | Bresenham 画线 | 20 |
|
||||
| `set_pixel(pixels, w, h, point, color)` | 设置单像素(带边界检查) | 8 |
|
||||
| `random_point(rng)` | 返回画布内随机点 `(x, y)` | 5 |
|
||||
| `stbi_write_png_to_func` 的回调封装 | 将 stb 输出收集到 `vector<uint8_t>` | 10 |
|
||||
|
||||
**FreeType 字体回退策略**(与 Pillow 行为对齐):
|
||||
|
||||
```
|
||||
1. 尝试 FT_New_Face(font_path) → 成功则使用
|
||||
2. 失败 → 尝试 DejaVuSans-Bold.ttf 默认路径
|
||||
3. 失败 → 打印错误到 stderr,退出进程(验证码无字体无法工作)
|
||||
(Python 版 Pillow 回退到 load_default() 位图字体无抗锯齿视觉效果差;
|
||||
C++ 版选择快速失败,在部署阶段确保字体安装)
|
||||
```
|
||||
|
||||
### 4.4 RpcServer 路由注册
|
||||
|
||||
```cpp
|
||||
int main(int argc, char* argv[]) {
|
||||
// 1. 命令行解析(boost::program_options)
|
||||
// 2. 配置加载:confs captcha.ini → 环境变量覆盖 → 完整性校验
|
||||
// 3. 初始化 CaptchaStore、字体路径
|
||||
// 4. 初始化日志客户端(AsyncQueue → db Logs 表)
|
||||
|
||||
pmc::net::RpcServer rpc(config.host, config.port,
|
||||
{.workers = config.workers});
|
||||
|
||||
CaptchaStore store;
|
||||
|
||||
// 路由
|
||||
rpc.add_route("image", [&](const json::object&) -> json::value {
|
||||
// 速率限制检查(审计 NEW-2)
|
||||
std::string text = generate_text(config.len);
|
||||
std::string token = store.generate(text, config.ttl, config.max_store);
|
||||
auto png = generate_image(text, config.font_path);
|
||||
json::object data;
|
||||
data["token"] = token;
|
||||
data["image"] = base64_encode(png);
|
||||
return data;
|
||||
});
|
||||
|
||||
rpc.add_route("verify", [&](const json::object& req) -> json::value {
|
||||
std::string token = json::value_to<std::string>(req.at("token"));
|
||||
std::string answer = json::value_to<std::string>(req.at("answer"));
|
||||
// 小写化
|
||||
std::transform(answer.begin(), answer.end(), answer.begin(), ::tolower);
|
||||
bool valid = store.verify(token, answer, config.ttl);
|
||||
json::object data;
|
||||
data["valid"] = valid;
|
||||
return data;
|
||||
});
|
||||
|
||||
rpc.add_route("health", [](const json::object&) -> json::value {
|
||||
json::object data;
|
||||
data["status"] = "ok";
|
||||
data["service"] = "captcha";
|
||||
return data;
|
||||
});
|
||||
|
||||
logs.info("captcha", "开始监听 tcp://" + config.host + ":" + std::to_string(config.port));
|
||||
rpc.start(); // 阻塞
|
||||
}
|
||||
```
|
||||
|
||||
### 4.5 LogC.hpp — 简化版异步日志客户端
|
||||
|
||||
参照 gateway `LogC.hpp`,基于 libmsg `AsyncQueue<LogEntry>`:
|
||||
|
||||
```cpp
|
||||
// ~60 行,与 log_client.py 功能等价
|
||||
// - 异步批量推送(攒满 50 条或 2 秒 flush)
|
||||
// - fail-safe(db 不可用静默丢弃)
|
||||
// - push()/info()/warning()/error() API 不变
|
||||
// - 环境变量 DB_HOST/DB_PORT 覆盖 db 地址
|
||||
```
|
||||
|
||||
### 4.6 速率限制(审计 NEW-2 修复)
|
||||
|
||||
在 `image` 路由中增加滑动窗口限流:
|
||||
|
||||
```cpp
|
||||
class RateLimiter {
|
||||
public:
|
||||
// 返回 true 表示允许,false 表示超限
|
||||
bool allow(int max_per_minute = 60) {
|
||||
std::lock_guard<std::mutex> lock(mtx_);
|
||||
auto now = std::chrono::steady_clock::now();
|
||||
|
||||
// 清理超过 60 秒的旧记录
|
||||
while (!timestamps_.empty() &&
|
||||
std::chrono::duration_cast<std::chrono::seconds>(
|
||||
now - timestamps_.front()).count() > 60) {
|
||||
timestamps_.pop_front();
|
||||
}
|
||||
|
||||
if ((int)timestamps_.size() >= max_per_minute)
|
||||
return false;
|
||||
|
||||
timestamps_.push_back(now);
|
||||
return true;
|
||||
}
|
||||
|
||||
private:
|
||||
std::deque<std::chrono::steady_clock::time_point> timestamps_;
|
||||
std::mutex mtx_;
|
||||
};
|
||||
```
|
||||
|
||||
通过在 `_handle_image` 入口处调用 `rate_limiter.allow()` 实现,超限时返回:
|
||||
|
||||
```json
|
||||
{"success": false, "error": "rate limit exceeded, please try again later"}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 5. 实现计划
|
||||
|
||||
### 5.1 分阶段交付
|
||||
|
||||
| 阶段 | 里程碑 | 关键产出 | 预估工时 |
|
||||
|------|--------|----------|----------|
|
||||
| **P0 — 基础设施** | 目录骨架 + Makefile + stb 引入 + LogC.hpp | `captcha2/` 结构搭建,能编译出空 `main()` 链接 libmsg + FreeType | 1h |
|
||||
| **P1 — 图片生成核心** | `generate_image()` 全链路 | FreeType 渲染 + 旋转 + 合成 + 干扰线 + 噪点 + PNG 编码,视觉效果与 Python 版对比 | 5h |
|
||||
| **P2 — CaptchaStore** | Token 存储 + 过期清理 + 容量保护 | 线程安全 `generate()` / `verify()`,单元测试 | 1.5h |
|
||||
| **P3 — RPC 服务集成** | `RpcServer` 三路由 + 命令行参数 | 能与 gateway 联调,pymsg 协议互通 | 2h |
|
||||
| **P4 — 配置加载** | confs 拉取 + env 覆盖 + 配置完整性校验 | 与现有 `captcha.ini` 完全兼容 | 1.5h |
|
||||
| **P5 — 安全加固** | 速率限制 + `mt19937` 安全种子 + 审计遗留修复 | 审计报告 NEW-1/NEW-2 全部修复 | 1h |
|
||||
| **P6 — 日志集成** | LogC.hpp 异步推送日志到 db | 结构化日志(生成/校验/错误),与 Python 版日志格式对齐 | 1h |
|
||||
| **P7 — 测试** | 黑盒测试 + 图片对比 + 压力测试 | `tests/test_captcha2.py` | 2.5h |
|
||||
| **P8 — 文档 + 部署** | README + _set.sh/_run.sh + init.txt 更新 | 完整文档 + 一键安装 + core 自启动 | 0.5h |
|
||||
|
||||
**总计预估:~16 小时**
|
||||
|
||||
### 5.2 P1 图片生成是高风险路径
|
||||
|
||||
P1 阶段需要逐像素对比 Python/C++ 版本的输出,验证视觉效果一致。策略:
|
||||
|
||||
1. **单元级对比**:对同一输入文本,分别运行 Python 和 C++ 版 `generate_image()`,保存 PNG 文件进行人工对比
|
||||
2. **关键差异容忍**:由于 `random` vs `mt19937` 的伪随机序列不同、FreeType 与 Pillow 的字体 hinting 差异,逐像素精确匹配不可能。目标是**同等干扰级别、同等人类可读性、同等机器识别难度**
|
||||
3. **灰度验证**:将两张图片都转灰度后,对比字符位置分布、旋转角度范围、噪点密度等统计特征
|
||||
|
||||
### 5.3 构建系统
|
||||
|
||||
```makefile
|
||||
# captcha2/Makefile
|
||||
CXX = g++
|
||||
CXXFLAGS = -std=c++20 -Wall -O2
|
||||
INCS = -I./inc -I/usr/include/freetype2
|
||||
LIBS = -lus_msg -lboost_json -lboost_program_options -lfreetype -lpthread
|
||||
|
||||
TARGET = captcha2.bin
|
||||
SRC = src/main.cpp
|
||||
|
||||
$(TARGET): $(SRC)
|
||||
$(CXX) $(CXXFLAGS) $(INCS) $^ $(LIBS) -o $@
|
||||
|
||||
install: $(TARGET)
|
||||
cp $(TARGET) /usr/local/bin/
|
||||
|
||||
clean:
|
||||
rm -f $(TARGET)
|
||||
```
|
||||
|
||||
> **说明**:stb_image_write.h 是单头文件(header-only),无需编译链接,直接 `#include` 并在一个翻译单元中 `#define STB_IMAGE_WRITE_IMPLEMENTATION`。FreeType 需链接 `-lfreetype`。
|
||||
|
||||
---
|
||||
|
||||
## 6. 测试计划
|
||||
|
||||
### 6.1 黑盒测试(`tests/test_captcha2.py`)
|
||||
|
||||
复用现有测试框架(`tests/blackbox/helpers.py` 中的 `PymsgClient`),覆盖:
|
||||
|
||||
| 测试用例 | 说明 |
|
||||
|----------|------|
|
||||
| `test_image_generation` | 调 `image` 路由,验证返回合法 token + base64 PNG(解码后为有效 PNG,尺寸 160×50) |
|
||||
| `test_verify_correct` | `image` → `verify(correct_answer)` → `{"valid": true}` |
|
||||
| `test_verify_wrong` | `image` → `verify(wrong_answer)` → `{"valid": false}` |
|
||||
| `test_verify_replay` | `image` → `verify(correct)` → `verify(correct)` → 第二次应返回 `invalid token`(一次性消费) |
|
||||
| `test_verify_expired` | 生成后等待 CAPTCHA_TTL+1 秒再校验 → `invalid token` |
|
||||
| `test_verify_no_token` | `verify` 缺少 token 字段 → `{"valid": false, "error": "token and answer required"}` |
|
||||
| `test_health` | 调 `health` 路由 → `{"status": "ok", "service": "captcha"}` |
|
||||
| `test_max_store` | 连续生成超过 `max_store` 个 → 后续返回 `too many pending captchas` |
|
||||
| `test_rate_limit` | 连续生成超过 `rate_limit_per_min` 次 → 返回 `rate limit exceeded`(与 Python 版的区别:Python 版暂无限流) |
|
||||
|
||||
### 6.2 图片质量验证
|
||||
|
||||
| 方法 | 说明 |
|
||||
|------|------|
|
||||
| 人工 A/B 对比 | 同一文本(固定随机种子),分别用 Python 和 C++ 生成,人工对比 |
|
||||
| 统计特征 | 字符分布范围、旋转角度范围、干扰线数量、噪点密度 → 与 Python 版统计对齐 |
|
||||
| OCR 挑战 | 用简单 OCR 工具(如 tesseract)分别尝试识别两版验证码,识别率应相近 |
|
||||
|
||||
### 6.3 gateway 联调
|
||||
|
||||
1. 启动 captcha2 服务替代 Python captcha
|
||||
2. `GET /captcha/image` → 浏览器正常显示验证码图片
|
||||
3. 登录/注册流程完整走通
|
||||
4. `X-Captcha-Token` 响应头正常
|
||||
|
||||
---
|
||||
|
||||
## 7. 迁移计划
|
||||
|
||||
### 7.1 渐进切换
|
||||
|
||||
```
|
||||
阶段 1: captcha2.bin 开发完成,通过全部测试
|
||||
↓
|
||||
阶段 2: 在测试环境与 Python captcha 并行运行(不同端口),验证
|
||||
↓
|
||||
阶段 3: 修改 core/init.txt,将 captcha/_run.sh 替换为 captcha2/_run.sh
|
||||
↓
|
||||
阶段 4: 生产切换,观察 1 周无异常后删除 Python captcha 目录
|
||||
```
|
||||
|
||||
### 7.2 gateway 零改动保证
|
||||
|
||||
- pymsg 协议不变(JSON RPC 约定不变)
|
||||
- 监听地址不变(`127.0.0.1:8009`)
|
||||
- `captcha.ini` 配置不变
|
||||
- gateway 的 `verify_captcha_remote()` 和 `handleCaptchaImage()` 通过 pymsg `MsgClient` 调用,与 Python 版字节级兼容
|
||||
|
||||
### 7.3 回滚方案
|
||||
|
||||
- `init.txt` 中恢复 `bash /opt/micro_servers/captcha/_run.sh`
|
||||
- Python venv 保持不变(不删除),可立即回滚
|
||||
|
||||
---
|
||||
|
||||
## 8. 安全考量
|
||||
|
||||
### 8.1 审计遗留修复
|
||||
|
||||
| 编号 | 问题 | 修复方式 |
|
||||
|------|------|----------|
|
||||
| NEW-1 | `generate_text()` 使用 `random.choices` 非安全随机 | C++ 版使用 `std::mt19937` 以 `std::random_device` 种子,Token 使用 256-bit `random_device` 直接生成 |
|
||||
| NEW-2 | 无生成速率限制,可被用于 CPU 资源消耗攻击 | 新增滑动窗口限流(默认 60 次/分钟,可配 `rate_limit_per_min`) |
|
||||
|
||||
### 8.2 已有安全机制保持
|
||||
|
||||
| 机制 | Python 版 | C++ 版 |
|
||||
|------|----------|--------|
|
||||
| Token 一次性使用 | `captcha_store.pop(token)` | `store_.erase(it)` 后返回 |
|
||||
| 5 分钟过期 | 惰性清理 | 惰性清理(与 Python 版一致) |
|
||||
| 常数时间比较 | XOR 逐字符 | XOR 逐字符(与 Python 版一致) |
|
||||
| 仅监听 127.0.0.1 | `RpcServer("127.0.0.1", ...)` | 同 |
|
||||
| 容量上限 | `MAX_STORE_SIZE = 10000` | 同 |
|
||||
| PNG base64 编码 | `base64.b64encode` | 手写或 libcrypt `Base64` |
|
||||
| 配置白名单合并 | `_merge_ini` 仅覆盖已知 key | `CaptchaConfig` 仅解析已知字段 |
|
||||
| 字体缺失回退 | Pillow `load_default()` | 尝试默认路径 → 失败则退出(快速失败优于静默降级为无抗锯齿位图字体) |
|
||||
|
||||
### 8.3 RpcServer 安全特性(继承自 libmsg)
|
||||
|
||||
- JSON 解析嵌套深度限制 32 层
|
||||
- 错误信息脱敏(不泄露内部路径/类型名)
|
||||
- `recv_timeout` 防 Slowloris(默认 10s)
|
||||
- Handler 异常不导致服务崩溃
|
||||
|
||||
---
|
||||
|
||||
## 9. 新增文件清单
|
||||
|
||||
```
|
||||
captcha2/
|
||||
├── src/
|
||||
│ └── main.cpp # ~420 行,全部实现
|
||||
├── inc/
|
||||
│ ├── stb_image_write.h # 公共领域单头 PNG 编码器(直接引入)
|
||||
│ └── LogC.hpp # ~60 行,简化版异步日志客户端
|
||||
├── Makefile # ~30 行
|
||||
├── _set.sh # make + install(~10 行)
|
||||
├── _run.sh # 启动脚本(~5 行)
|
||||
└── README.md # 模块文档
|
||||
```
|
||||
|
||||
**0 个新增文件需修改现有代码。** gateway / core / waf / init.txt 均零改动。
|
||||
|
||||
---
|
||||
|
||||
## 10. 依赖清单
|
||||
|
||||
| 依赖 | 用途 | 是否新增 | 安装方式 |
|
||||
|------|------|----------|----------|
|
||||
| `libus_msg.so` (libmsg) | RPC 服务端 + 异步日志 | 否 | 已通过 `_inst.sh` 安装 |
|
||||
| `libboost_json` | JSON 解析 | 否 | 已通过 `_setup.deb.sh` 安装 |
|
||||
| `libboost_program_options` | 命令行参数 | 否 | 已通过 `_setup.deb.sh` 安装 |
|
||||
| `libfreetype-dev` | TrueType 字体渲染 | **是** | `apt install libfreetype-dev`(已在 `_setup.deb.sh` 中添加) |
|
||||
| `stb_image_write.h` | PNG 编码 | 形式上是 | 单头文件直接放入 `inc/`,无编译依赖 |
|
||||
|
||||
### 10.1 `_setup.deb.sh` 新增依赖
|
||||
|
||||
在根目录 `_setup.deb.sh` 的 `apt install` 命令中追加 `libfreetype-dev`:
|
||||
|
||||
```bash
|
||||
sudo apt install -y ... libfreetype-dev
|
||||
```
|
||||
|
||||
该包在 Ubuntu 24.04 中预装,大小约 1.2MB(含 `.so`)。
|
||||
|
||||
---
|
||||
|
||||
## 11. 与 Python 版的差异汇总
|
||||
|
||||
| 项目 | Python 版 | C++ 版 | 影响 |
|
||||
|------|----------|--------|------|
|
||||
| 运行时 | Python 3 + venv + Pillow | 单二进制 `captcha2.bin` | 消除 Python 依赖 |
|
||||
| 随机数 | `random.choices` + `secrets.token_hex(16)` | `mt19937(random_device)` + 256-bit random_device | 安全增强 |
|
||||
| 速率限制 | 无 | 滑动窗口 60 次/分钟(可配) | 审计 NEW-2 修复 |
|
||||
| 字体回退 | Pillow `load_default()`(无抗锯齿位图字体) | 快速失败(打印错误退出) | 部署时确保字体安装 |
|
||||
| Token 长度 | 128-bit hex (32 字符) | 256-bit hex (64 字符) | 对调用方透明(Token 只做 key 用) |
|
||||
| log_client | Python `AsyncRpcQueue` | C++ `AsyncQueue<LogEntry>` | 功能等价 |
|
||||
| 进程模型 | Python GIL | 原生多线程(`workers` 配置) | 真实并发图片生成 |
|
||||
|
||||
---
|
||||
|
||||
## 12. 风险与应对
|
||||
|
||||
| 风险 | 概率 | 影响 | 应对 |
|
||||
|------|------|------|------|
|
||||
| FreeType 与 Pillow 字体渲染差异大,验证码视觉效果不一致 | 中 | 中 | P1 阶段人工对比 + 统计特征验证;必要时调整旋转/位置参数 |
|
||||
| 像素旋转合成逻辑有 bug,字符变形/丢失 | 低 | 高 | 大量单元测试 + 固定种子生成保存图片对比 |
|
||||
| FreeType 字体缺失导致进程退出 | 低 | 高 | `_setup.deb.sh` 确保 `fonts-dejavu-core` 安装;启动时显式检查字体文件存在 |
|
||||
| gateway 联调时 pymsg 协议兼容性问题 | 极低 | 高 | libmsg 已在 db2/dnfs2/fsys 中经过充分验证;协议极简 |
|
||||
| `max_store` 容量保护逻辑与 Python 版不一致 | 低 | 低 | 黑盒测试覆盖边界条件 |
|
||||
|
||||
---
|
||||
|
||||
## A. 附录:stb_image_write.h 许可说明
|
||||
|
||||
`stb_image_write.h` 是 Sean Barrett 编写的单头文件公共领域库(public domain / MIT),广泛用于游戏引擎和图形工具。将其放入 `captcha2/inc/` 目录即可使用,无需额外安装。文件约 1800 行,提供:
|
||||
|
||||
```c
|
||||
// 主要 API
|
||||
STBIWDEF int stbi_write_png(char const *filename, int w, int h, int comp, const void *data, int stride);
|
||||
STBIWDEF int stbi_write_png_to_func(stbi_write_func *func, void *context, int w, int h, int comp, const void *data, int stride);
|
||||
```
|
||||
|
||||
本项目使用 `stbi_write_png_to_func`(内存输出),在 `main.cpp` 中:
|
||||
|
||||
```cpp
|
||||
#define STB_IMAGE_WRITE_IMPLEMENTATION
|
||||
#include "stb_image_write.h"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 13. 实现后记 · 审计与修复
|
||||
|
||||
### 13.1 实现完成
|
||||
|
||||
2026-08-09,captcha2 按本方案完成全部开发(P0~P6),编译为 434KB 动态链接 ELF 二进制。
|
||||
目录结构:
|
||||
|
||||
```
|
||||
captcha2/
|
||||
├── Makefile # g++ -std=c++20 -O2
|
||||
├── _run.sh / _set.sh # 启动/编译脚本
|
||||
├── inc/
|
||||
│ ├── LogC.hpp # 异步日志客户端(AsyncQueue → db2:8005)
|
||||
│ └── stb_image_write.h # PNG 编码(1724 行,public domain)
|
||||
├── src/
|
||||
│ └── main.cpp # 主程序(672 行)
|
||||
├── test_gen.cpp # 图片生成独立测试
|
||||
├── test_rpc.py # pymsg RPC 全功能测试(7 项自动化)
|
||||
└── captcha2.bin # 编译产物
|
||||
```
|
||||
|
||||
### 13.2 安全审计
|
||||
|
||||
同日完成全量安全审计(参见 `docs/security/captcha2_SECURITY_AUDIT_20260809.md`),
|
||||
发现 11 项问题(P0×1 / P1×1 / P2×3 / P3×6),其中:
|
||||
|
||||
| 编号 | 等级 | 问题 | 状态 |
|
||||
|------|------|------|------|
|
||||
| P0-1 | 🔴 | 存储答案未小写化,含大写字母的验证码约 90% 永远校验失败 | ✅ 已修复 |
|
||||
| P2-2 | 🟡 | `::tolower(char)` 对非 ASCII 输入为 UB | ✅ 已修复 |
|
||||
| P1-1 | 🟠 | 全局速率限制可被单一匿名用户 DoS | ❌ 待修复(需 gateway 层按 IP 限流) |
|
||||
| P2-1 | 🟡 | 答案生成用 mt19937 非密码学安全随机 | ❌ 遗留(与 Python 版同源) |
|
||||
|
||||
### 13.3 P0-1 修复详情
|
||||
|
||||
**根因**: Python 原版存储时即小写化答案(`text.lower()`),C++ 初版直接存原始混合大小写文本,
|
||||
verify 路由仅对用户输入做 lowercase,导致大小写不匹配。
|
||||
|
||||
**修复** (`src/main.cpp` L588-597):
|
||||
```cpp
|
||||
std::string text = generate_text(g_config.len);
|
||||
// 小写化存储,与 Python 版行为对齐(大小写不敏感校验)
|
||||
std::string answer_lc = text;
|
||||
std::transform(answer_lc.begin(), answer_lc.end(), answer_lc.begin(),
|
||||
[](unsigned char c) { return static_cast<char>(std::tolower(c)); });
|
||||
token = g_store.generate(answer_lc, g_config.ttl, g_config.max_store);
|
||||
```
|
||||
|
||||
同时修复 `::tolower` UB(L633-636),改用 `unsigned char` lambda 包装。
|
||||
|
||||
### 13.4 init.txt 集成
|
||||
|
||||
`core/init.txt` 已将 `captcha/_run.sh` 替换为 `captcha2/_run.sh`;
|
||||
`confs/configs/captcha.ini` 已补充 `max_store=10000` 和 `rate_limit_per_min=60` 字段。
|
||||
@@ -1,740 +0,0 @@
|
||||
# confs 模块 C++20 重写开发方案
|
||||
|
||||
> 文档类型:功能开发计划
|
||||
> 模块:confs → confs2(C++20 重写)
|
||||
> 作者:自动生成
|
||||
> 日期:2026-08-09
|
||||
|
||||
---
|
||||
|
||||
## 1. 概述
|
||||
|
||||
### 1.1 背景与动机
|
||||
|
||||
`confs` 统一配置服务当前以 Python 实现,基于 `pymsg RpcServer` 在 `127.0.0.1:8010` 提供配置文件集中托管与分发。5 个内部服务(static、dnfs、db、captcha、redirect80)通过 `pymsg RpcClient` 消费配置。
|
||||
|
||||
本项目已完成 db2、captcha2、smtp3、dnfs2 等 C++20 重写,遵循"原版与重写版并存 → 重写版验证通过后切换"的渐进策略。confs 作为最后一个纯 Python 基础设施服务,其 C++ 重写可:
|
||||
|
||||
| 收益 | 说明 |
|
||||
|------|------|
|
||||
| 内存占用降低 | Python 进程 ~50-100MB RSS → C++ 预估 < 5MB |
|
||||
| 启动速度提升 | Python 解释器启动 ~200ms → C++ 瞬时 |
|
||||
| 依赖链简化 | 消除 Python venv / pip / nuitka 工具链 |
|
||||
| 安全审计遗留修复 | 重写时一并修复 Python 版 7 项遗留审计发现 |
|
||||
| 未来合入 core 的基础 | C++ 同构后,未来可评估合入 core 注册中心 |
|
||||
|
||||
### 1.2 设计目标
|
||||
|
||||
1. **协议零改动**:pymsg RPC 协议、端口 `8010`、路由名(`get`/`list`/`reload`/`health`)、响应格式全部不变,5 个消费方零改动。
|
||||
2. **功能完全对等**:ConfigStore 内存缓存 + mtime 自动重载、INI/JSON 按键取值、扩展名白名单、TOCTOU 防护、文件大小限制——逐项 C++ 实现。
|
||||
3. **零新增系统依赖**:仅链接 `libmsg`(`-lus_msg`)、`boost_json`、`boost_program_options`、`pthread`——均为系统已安装库。
|
||||
4. **安全审计遗留修复**:Python 版 7 项未修复发现(P2-3/P2-4/P3-1~P3-6)在重写中一并解决。
|
||||
5. **配置目录兼容**:复用现有 `confs/configs/` 目录,配置文件零迁移。
|
||||
6. **可演进**:保留向 core 注册中心注册的接口预留,未来合入评估时改动最小。
|
||||
|
||||
### 1.3 不做什么
|
||||
|
||||
- **不改协议**:不引入 HTTP,不改变 pymsg 帧格式,不增减路由。
|
||||
- **不改消费方**:5 个 Python 服务的 `RpcClient` 调用代码不改。
|
||||
- **不引入新解析库**:INI 解析手写(与 captcha2 的 `merge_ini` 风格一致),不引入 `boost::property_tree::ini_parser`(其大小写语义与 Python ConfigParser 不一致)。
|
||||
- **不实现 YAML/TOML 按键取值**:与 Python 版一致,仅返回全文。
|
||||
|
||||
---
|
||||
|
||||
## 2. 现状分析
|
||||
|
||||
### 2.1 Python confs 源码结构
|
||||
|
||||
```
|
||||
confs/
|
||||
├── confs.py # 服务主程序(ConfigStore + ConfsServer,361 行)
|
||||
├── log_client.py # 异步日志客户端 SDK(160 行)
|
||||
├── configs/ # 配置文件存放目录(9 个 .ini 文件)
|
||||
│ ├── ca.ini
|
||||
│ ├── captcha.ini
|
||||
│ ├── db.ini
|
||||
│ ├── dnfs.ini
|
||||
│ ├── example.ini
|
||||
│ ├── gateway.ini
|
||||
│ ├── redirect80.ini
|
||||
│ ├── smtp.ini
|
||||
│ └── static.ini
|
||||
├── _run.sh # 启动脚本(venv 中 exec python3 confs.py)
|
||||
├── _set.sh # 初始化脚本(建 venv + pip install)
|
||||
├── requirements.txt # 依赖(pymsg + nuitka)
|
||||
└── README.md
|
||||
```
|
||||
|
||||
### 2.2 关键组件
|
||||
|
||||
| 组件 | 职责 | 行数 |
|
||||
|------|------|------|
|
||||
| `ConfigStore` | 磁盘 configs/ 目录 + 内存缓存 + mtime 自动重载 + 路径安全 | ~120 行 |
|
||||
| `ConfsServer` | RpcServer 入口,路由分发(get/list/reload/health) | ~90 行 |
|
||||
| `_extract_key()` | INI `Section.option` 取值 + JSON 点路径取值 | ~40 行 |
|
||||
| `log_client.py` | 异步批量推送日志到 db:8005 | 160 行 |
|
||||
|
||||
### 2.3 现有路由
|
||||
|
||||
| 路由 | 方法 | 说明 |
|
||||
|------|------|------|
|
||||
| `get` | `{"name": "db.ini"}` | 取配置全文;带 `key` 时返回单键值(仅 ini/json) |
|
||||
| `list` | `{}` | 列出配置目录全部文件及元信息 |
|
||||
| `reload` | `{"name": "db.ini"}` 或 `{}` | 重载指定配置;省略 name 清空全部缓存 |
|
||||
| `health` | `{}` | 健康检查 |
|
||||
|
||||
### 2.4 消费方清单
|
||||
|
||||
| 服务 | 配置文件 | 接入方式 |
|
||||
|------|---------|---------|
|
||||
| static | `static.ini` | pymsg `RpcClient` + `retry_on_refused=True` |
|
||||
| dnfs | `dnfs.ini` | pymsg `RpcClient` + `retry_on_refused=True` |
|
||||
| db | `db.ini` | pymsg `RpcClient` + `retry_on_refused=True` |
|
||||
| captcha | `captcha.ini` | pymsg `RpcClient` + `retry_on_refused=True` |
|
||||
| redirect80 | `redirect80.ini` | pymsg `RpcClient` + `retry_on_refused=True` |
|
||||
|
||||
### 2.5 Python 版安全审计遗留(7 项)
|
||||
|
||||
| 编号 | 级别 | 问题 | 重写修复方案 |
|
||||
|------|------|------|-------------|
|
||||
| P2-3 | P2 | 配置名/key 直接拼入日志,未过滤控制字符 | 日志消息过滤 `\n`/`\r`/ANSI 转义 |
|
||||
| P2-4 | P2 | `list()` 持锁遍历目录+stat | 拷贝文件名列表后释放锁再 stat |
|
||||
| P3-1 | P3 | `_run.sh` 有 `sleep` hack | 移除 sleep,靠 init.txt 启动顺序 |
|
||||
| P3-2 | P3 | `_set.sh` 硬编码 pip 镜像源 | C++ 重写无需 pip,`_set.sh` 改为 `make` |
|
||||
| P3-3 | P3 | `_set.sh` 缺少 `set -e` | `_set.sh` 加 `set -euo pipefail` |
|
||||
| P3-4 | P3 | requirements.txt 含运行时不需要的 nuitka | C++ 无 requirements.txt |
|
||||
| P3-5 | P3 | ConfigParser 大小写隐式行为未文档化 | INI 解析器明确定义大小写策略 |
|
||||
| P3-6 | P3 | 错误消息暴露底层 IO 异常 | 异常消息脱敏,不返回系统路径 |
|
||||
|
||||
---
|
||||
|
||||
## 3. 整体设计
|
||||
|
||||
### 3.1 架构总览
|
||||
|
||||
```
|
||||
┌──────────────────────────────────────────────────────────┐
|
||||
│ confs2 (C++20, :8010) │
|
||||
│ │
|
||||
│ ┌────────────────────────────────────────────────────┐ │
|
||||
│ │ ConfigStore (线程安全) │ │
|
||||
│ │ ├── 磁盘: configs/ 目录 │ │
|
||||
│ │ ├── 内存缓存: {name → {content, mtime, size}} │ │
|
||||
│ │ ├── mtime 自动重载 │ │
|
||||
│ │ └── 安全: 路径遍历/TOCTOU/扩展名白名单/大小限制 │ │
|
||||
│ └───────────────────────┬────────────────────────────┘ │
|
||||
│ │ │
|
||||
│ ┌───────────────────────▼────────────────────────────┐ │
|
||||
│ │ RpcServer (libmsg, :8010, pymsg 协议) │ │
|
||||
│ │ ├── get → ConfigStore.get() + _extract_key() │ │
|
||||
│ │ ├── list → ConfigStore.list() │ │
|
||||
│ │ ├── reload → ConfigStore.reload() │ │
|
||||
│ │ └── health → 静态响应 │ │
|
||||
│ └────────────────────────────────────────────────────┘ │
|
||||
│ │
|
||||
│ ┌────────────────────────────────────────────────────┐ │
|
||||
│ │ LogC (异步日志, libmsg AsyncQueue → db:8005) │ │
|
||||
│ └────────────────────────────────────────────────────┘ │
|
||||
└──────────────────────────────────────────────────────────┘
|
||||
▲
|
||||
│ pymsg RPC (TCP, 127.0.0.1:8010)
|
||||
│
|
||||
┌─────┴─────┬──────┬──────┬──────┬──────┐
|
||||
▼ ▼ ▼ ▼ ▼ ▼
|
||||
static dnfs db captcha redirect80
|
||||
(RpcClient, 消费方零改动)
|
||||
```
|
||||
|
||||
### 3.2 核心组件
|
||||
|
||||
| 组件 | 职责 | 对应 Python |
|
||||
|------|------|-------------|
|
||||
| **ConfigStore** | 配置文件仓库(磁盘+缓存+安全) | `ConfigStore` 类 |
|
||||
| **IniParser** | INI 解析(Section.option 取值) | `configparser.ConfigParser` |
|
||||
| **JsonKeyExtractor** | JSON 点路径取值 | `_extract_key` JSON 分支 |
|
||||
| **LogC** | 异步日志推送 | `log_client.py` |
|
||||
| **RpcServer** | pymsg RPC 服务端(libmsg) | `pymsg.RpcServer` |
|
||||
|
||||
### 3.3 目录结构
|
||||
|
||||
```
|
||||
confs2/
|
||||
├── src/
|
||||
│ └── main.cpp # 全部源码(ConfigStore + 路由 + main)
|
||||
├── inc/
|
||||
│ └── LogC.hpp # 异步日志客户端(从 captcha2 改编)
|
||||
├── configs/ # → 符号链接到 ../confs/configs/(配置文件零迁移)
|
||||
├── Makefile # 构建系统(g++ -std=c++20)
|
||||
├── _run.sh # 启动脚本
|
||||
├── _set.sh # 编译脚本
|
||||
└── README.md # 模块文档
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 4. 详细设计
|
||||
|
||||
### 4.1 ConfigStore — 配置文件仓库
|
||||
|
||||
```cpp
|
||||
// 核心数据结构
|
||||
struct CacheEntry {
|
||||
std::string content;
|
||||
std::filesystem::file_time_type mtime;
|
||||
std::size_t size;
|
||||
};
|
||||
|
||||
class ConfigStore {
|
||||
public:
|
||||
explicit ConfigStore(const std::string& root);
|
||||
|
||||
// 取配置全文(带缓存 + mtime 自动重载)
|
||||
// 返回 CacheEntry 的拷贝(防缓存污染,与 Python 版 dict(entry) 语义一致)
|
||||
CacheEntry get(const std::string& name);
|
||||
|
||||
// 列出配置目录全部文件元信息
|
||||
std::vector<ConfigMeta> list();
|
||||
|
||||
// 重载指定配置 / 清空全部缓存
|
||||
std::vector<std::string> reload(const std::string* name);
|
||||
|
||||
std::size_t cached_count() const;
|
||||
|
||||
private:
|
||||
// 扩展名白名单校验(纯字符串操作,不访问文件系统)
|
||||
static void validate_name(const std::string& name);
|
||||
|
||||
// 持锁解析路径(realpath 前缀校验 + 文件存在性)
|
||||
std::string resolve_locked(const std::string& name);
|
||||
|
||||
std::filesystem::path root_;
|
||||
std::unordered_map<std::string, CacheEntry> cache_;
|
||||
mutable std::mutex mtx_;
|
||||
};
|
||||
```
|
||||
|
||||
#### 安全机制(逐项 C++ 实现)
|
||||
|
||||
| 安全机制 | Python 实现 | C++ 实现 |
|
||||
|----------|-------------|----------|
|
||||
| 扩展名白名单 | `ALLOWED_EXT = {".ini", ...}` | `static const std::unordered_set<std::string> kAllowedExt` |
|
||||
| 路径遍历防护 | `os.path.basename` + `os.path.realpath` 前缀校验 | `std::filesystem::path::filename()` + `canonical()` 前缀校验 |
|
||||
| TOCTOU 防护 | 锁内 `realpath→stat→open` | 锁内 `canonical→file_size→ifstream` |
|
||||
| 文件大小限制 | `MAX_CONFIG_SIZE = 1MiB`,stat 后校验 | `file_size()` 后校验 `kMaxConfigSize` |
|
||||
| 缓存防污染 | 返回 `dict(entry)` 浅拷贝 | 返回 `CacheEntry` 值拷贝 |
|
||||
| mtime 自动重载 | `os.stat().st_mtime` 比对 | `last_write_time()` 比对 |
|
||||
| list 符号链接防护 | 遍历时 `realpath` 前缀校验 | 遍历时 `canonical()` 前缀校验,逃逸链接跳过 |
|
||||
| 日志注入防护 | **未实现(P2-3 遗留)** | **新增**:日志消息过滤 `\n`/`\r`/`\x1b` |
|
||||
|
||||
#### TOCTOU 防护流程(锁内原子操作)
|
||||
|
||||
```cpp
|
||||
CacheEntry ConfigStore::get(const std::string& name) {
|
||||
validate_name(name); // 锁外快速拒绝非法名
|
||||
std::lock_guard<std::mutex> lock(mtx_);
|
||||
|
||||
auto path = resolve_locked(name); // 锁内: canonical() 解析符号链接 + 前缀校验
|
||||
|
||||
// 锁内: stat → mtime 比对 → open → read,无 TOCTOU 窗口
|
||||
auto mtime = std::filesystem::last_write_time(path);
|
||||
auto size = std::filesystem::file_size(path);
|
||||
|
||||
auto it = cache_.find(name);
|
||||
if (it != cache_.end() && it->second.mtime == mtime)
|
||||
return it->second; // 缓存命中,返回拷贝
|
||||
|
||||
if (size > kMaxConfigSize)
|
||||
throw ConfigError("配置文件过大");
|
||||
|
||||
std::ifstream f(path);
|
||||
std::string content((std::istreambuf_iterator<char>(f)),
|
||||
std::istreambuf_iterator<char>());
|
||||
|
||||
CacheEntry entry{content, mtime, content.size()};
|
||||
cache_[name] = entry;
|
||||
return entry; // 返回拷贝
|
||||
}
|
||||
```
|
||||
|
||||
### 4.2 IniParser — INI 按键取值
|
||||
|
||||
与 Python `ConfigParser(interpolation=None)` 行为对等的 C++ 手写解析器:
|
||||
|
||||
| Python ConfigParser 行为 | C++ 实现策略 |
|
||||
|--------------------------|-------------|
|
||||
| Section 大小写敏感 | `std::string` 原样比较 |
|
||||
| Option 大小写不敏感 | 比较前 `std::tolower` 归一化 |
|
||||
| `;` 和 `#` 注释 | 行首匹配跳过 |
|
||||
| `=` 和 `:` 分隔符 | 仅支持 `=`(与现有 configs/ 一致) |
|
||||
| 禁用 `%(key)s` 插值 | 不实现插值,原样返回 |
|
||||
| 多行续行(前导空格) | 支持:下一行以空格/Tab 开头则追加 |
|
||||
|
||||
```cpp
|
||||
// 从 INI 文本中提取 Section.option 的值
|
||||
// 键格式: "Section.option"
|
||||
// 大小写: section 敏感, option 不敏感(与 Python 一致)
|
||||
// 插值: 禁用,原样返回(与 Python interpolation=None 一致)
|
||||
static std::optional<std::string> extract_ini_key(
|
||||
const std::string& content, const std::string& key);
|
||||
```
|
||||
|
||||
### 4.3 JsonKeyExtractor — JSON 点路径取值
|
||||
|
||||
```cpp
|
||||
// 从 JSON 文本中按点路径提取值
|
||||
// 键格式: "server.port" 或 "servers.0.host"(数字下标)
|
||||
// 返回 boost::json::value(可为 string/int/double/bool/null/array/object)
|
||||
static json::value extract_json_key(
|
||||
const std::string& content, const std::string& key);
|
||||
```
|
||||
|
||||
### 4.4 路由实现
|
||||
|
||||
#### `get` — 取配置
|
||||
|
||||
```cpp
|
||||
rpc.add_route("get", [](const json::object& req) -> json::value {
|
||||
// 参数校验
|
||||
auto name_it = req.find("name");
|
||||
if (name_it == req.end() || !name_it->value().is_string())
|
||||
throw ConfigError("缺少 name 字段");
|
||||
|
||||
std::string name = json::value_to<std::string>(name_it->value());
|
||||
auto entry = g_store.get(name);
|
||||
|
||||
// 无 key:返回全文
|
||||
auto key_it = req.find("key");
|
||||
if (key_it == req.end())
|
||||
return json::object{
|
||||
{"name", name},
|
||||
{"content", entry.content},
|
||||
{"size", (std::int64_t)entry.size},
|
||||
{"mtime", iso_timestamp(entry.mtime)}
|
||||
};
|
||||
|
||||
// 有 key:按键取值
|
||||
std::string key = json::value_to<std::string>(key_it->value());
|
||||
auto ext = get_extension(name);
|
||||
|
||||
if (ext == ".ini" || ext == ".conf" || ext == ".cfg") {
|
||||
auto val = extract_ini_key(entry.content, key);
|
||||
if (!val) throw KeyNotFound("键不存在: " + name + " -> " + key);
|
||||
return json::object{{"name", name}, {"key", key}, {"value", *val}};
|
||||
}
|
||||
if (ext == ".json") {
|
||||
auto val = extract_json_key(entry.content, key);
|
||||
return json::object{{"name", name}, {"key", key}, {"value", val}};
|
||||
}
|
||||
throw ConfigError(get_extension(name) + " 类型不支持按键查询");
|
||||
});
|
||||
```
|
||||
|
||||
#### `list` — 列出配置
|
||||
|
||||
```cpp
|
||||
rpc.add_route("list", [](const json::object&) -> json::value {
|
||||
auto configs = g_store.list();
|
||||
json::array arr;
|
||||
for (const auto& c : configs)
|
||||
arr.push_back(json::object{
|
||||
{"name", c.name}, {"size", (std::int64_t)c.size},
|
||||
{"mtime", c.mtime_iso}, {"cached", c.cached}
|
||||
});
|
||||
return json::object{
|
||||
{"config_dir", g_store.root_string()},
|
||||
{"configs", arr}
|
||||
};
|
||||
});
|
||||
```
|
||||
|
||||
#### `reload` — 重载配置
|
||||
|
||||
```cpp
|
||||
rpc.add_route("reload", [](const json::object& req) -> json::value {
|
||||
auto it = req.find("name");
|
||||
|
||||
// 省略 name 或 null → 清空全部缓存
|
||||
if (it == req.end() || it->value().is_null()) {
|
||||
auto names = g_store.reload(nullptr);
|
||||
return json::object{{"reloaded", make_json_array(names)}};
|
||||
}
|
||||
|
||||
// 非字符串 → 报错
|
||||
if (!it->value().is_string())
|
||||
throw ConfigError("name 字段必须是字符串");
|
||||
|
||||
std::string name = json::value_to<std::string>(it->value());
|
||||
// 空字符串 → 报错(不误清全部缓存,与 Python 版一致)
|
||||
if (name.empty())
|
||||
throw ConfigError("name 字段不能为空(要清空全部缓存请省略 name 字段)");
|
||||
|
||||
auto names = g_store.reload(&name);
|
||||
return json::object{{"reloaded", make_json_array(names)}};
|
||||
});
|
||||
```
|
||||
|
||||
#### `health` — 健康检查
|
||||
|
||||
```cpp
|
||||
rpc.add_route("health", [](const json::object&) -> json::value {
|
||||
return json::object{
|
||||
{"status", "ok"},
|
||||
{"service", "confs"},
|
||||
{"config_dir", g_store.root_string()},
|
||||
{"cached", (std::int64_t)g_store.cached_count()},
|
||||
{"uptime", uptime_seconds()}
|
||||
};
|
||||
});
|
||||
```
|
||||
|
||||
### 4.5 配置加载(confs 自身的引导配置)
|
||||
|
||||
confs 自身不需要从 confs 获取配置(避免循环依赖),使用硬编码默认值 + 环境变量覆盖 + 命令行参数覆盖:
|
||||
|
||||
```cpp
|
||||
struct ConfsConfig {
|
||||
std::string host = "127.0.0.1";
|
||||
int port = 8010;
|
||||
int workers = 4;
|
||||
std::string config_dir; // 默认: 同级 configs/ 目录
|
||||
};
|
||||
|
||||
// 加载顺序: 硬编码默认 → 环境变量覆盖 → 命令行覆盖
|
||||
// 环境变量: LISTEN_ADDR, LISTEN_PORT, WORKERS, CONFS_DIR
|
||||
```
|
||||
|
||||
### 4.6 日志客户端(LogC.hpp)
|
||||
|
||||
从 [liblogc/inc/us/LogC.hpp](../../liblogc/inc/us/LogC.hpp)(原 captcha2/inc/LogC.hpp,后独立为 liblogc 模块)改编,仅改 namespace 和服务名:
|
||||
|
||||
```cpp
|
||||
namespace confs2 {
|
||||
// ... 与 captcha2::LogC 结构完全一致
|
||||
// 仅改动:
|
||||
// - namespace: captcha2 → confs2
|
||||
// - 日志中服务名: "captcha2" → "confs"
|
||||
// - 客户端名: "captcha2-log" → "confs-log"
|
||||
}
|
||||
```
|
||||
|
||||
#### 日志注入防护(修复 P2-3)
|
||||
|
||||
```cpp
|
||||
// 过滤日志消息中的控制字符(\n, \r, \x1b ANSI 转义等)
|
||||
static std::string sanitize_for_log(std::string s) {
|
||||
s.erase(std::remove_if(s.begin(), s.end(), [](unsigned char c) {
|
||||
return c < 0x20 && c != '\t'; // 保留 Tab,过滤其他控制字符
|
||||
}), s.end());
|
||||
return s;
|
||||
}
|
||||
```
|
||||
|
||||
### 4.7 异常消息脱敏(修复 P3-6)
|
||||
|
||||
```cpp
|
||||
// 异常消息脱敏:不返回系统路径、权限错误等底层信息
|
||||
static std::string sanitize_error_msg(const std::string& msg) {
|
||||
// 替换绝对路径为 <redacted>
|
||||
static const std::regex path_re(R"(/\S+(?:/\S*)*)");
|
||||
std::string result = std::regex_replace(msg, path_re, "<redacted>");
|
||||
// 过滤 Permission denied 等系统错误
|
||||
static const std::regex perm_re(R"(Permission denied|No such file)");
|
||||
result = std::regex_replace(result, perm_re, "access error");
|
||||
return result;
|
||||
}
|
||||
```
|
||||
|
||||
> 注意:RpcServer 已内置 `_sanitize_error()` 方法(见 [RpcServer.hpp](../../libmsg/inc/us/RpcServer.hpp#L188-L197)),会自动替换绝对路径。confs2 的 handler 抛出的异常会先经 RpcServer 脱敏。但 confs2 仍应在业务层对底层 IO 异常做一层包装,将 `Permission denied` 等系统错误转为业务语义错误。
|
||||
|
||||
---
|
||||
|
||||
## 5. 实现计划
|
||||
|
||||
### 5.1 分阶段交付
|
||||
|
||||
| 阶段 | 里程碑 | 涉及文件 | 预估工时 |
|
||||
|------|--------|----------|----------|
|
||||
| **P0 — 骨架搭建** | Makefile + _run.sh + _set.sh + 空 main.cpp(RpcServer 启动 + health 路由) | 新增 `confs2/Makefile`、`_run.sh`、`_set.sh`、`src/main.cpp`、`inc/LogC.hpp` | 1h |
|
||||
| **P1 — ConfigStore 核心** | ConfigStore 类:磁盘读取 + 内存缓存 + mtime 重载 + 安全机制 | `src/main.cpp` | 3h |
|
||||
| **P2 — 路由实现** | get / list / reload / health 四路由完整实现 | `src/main.cpp` | 2h |
|
||||
| **P3 — INI 按键取值** | IniParser 手写 + JSON 点路径取值 | `src/main.cpp` | 2h |
|
||||
| **P4 — 日志客户端** | LogC.hpp 改编 + 日志注入防护 | `inc/LogC.hpp`、`src/main.cpp` | 1h |
|
||||
| **P5 — 配置加载** | 环境变量 + 命令行参数解析 | `src/main.cpp` | 1h |
|
||||
| **P6 — 编译验证** | 零警告编译 + ldd 依赖检查 + --help 验证 | `Makefile` | 0.5h |
|
||||
| **P7 — 功能测试** | 逐路由对比 Python 版响应一致性 | `tests/` 新增 | 2h |
|
||||
| **P8 — 切换 init.txt** | init.txt 改为 confs2/_run.sh + 生产验证 | `core/init.txt` | 0.5h |
|
||||
|
||||
**总计预估:~13 小时**
|
||||
|
||||
### 5.2 Makefile
|
||||
|
||||
```makefile
|
||||
CXX := g++
|
||||
CXXFLAGS := -std=c++20 -Wall -Wextra -O2
|
||||
INCS := -I./inc
|
||||
LIBS := -lus_msg -lboost_json -lboost_program_options -lpthread
|
||||
|
||||
TARGET := confs2.bin
|
||||
SRC := src/main.cpp
|
||||
|
||||
.PHONY: all clean install
|
||||
|
||||
all: $(TARGET)
|
||||
|
||||
$(TARGET): $(SRC)
|
||||
$(CXX) $(CXXFLAGS) $(INCS) $^ $(LIBS) -o $@
|
||||
|
||||
install: $(TARGET)
|
||||
cp $(TARGET) /usr/local/bin/
|
||||
|
||||
clean:
|
||||
rm -f $(TARGET)
|
||||
```
|
||||
|
||||
### 5.3 _run.sh
|
||||
|
||||
```bash
|
||||
#!/bin/bash
|
||||
cd "$(dirname "$0")" || exit
|
||||
sleep 1
|
||||
exec ./confs2.bin
|
||||
```
|
||||
|
||||
### 5.4 _set.sh
|
||||
|
||||
```bash
|
||||
#!/bin/bash
|
||||
set -euo pipefail
|
||||
cd "$(dirname "$0")" || exit
|
||||
|
||||
make clean 2>/dev/null || true
|
||||
make
|
||||
echo "[confs2] 编译完成 → ./confs2.bin"
|
||||
|
||||
# 创建 configs 符号链接(指向 ../confs/configs/,配置文件零迁移)
|
||||
if [ ! -e configs ]; then
|
||||
ln -s ../confs/configs configs
|
||||
echo "[confs2] 已创建 configs → ../confs/configs 符号链接"
|
||||
fi
|
||||
```
|
||||
|
||||
### 5.5 新增文件清单
|
||||
|
||||
```
|
||||
confs2/
|
||||
├── src/
|
||||
│ └── main.cpp # 全部业务源码(预估 ~500 行)
|
||||
├── inc/
|
||||
│ └── LogC.hpp # 异步日志客户端(从 captcha2 改编,~110 行)
|
||||
├── configs # → 符号链接到 ../confs/configs/
|
||||
├── Makefile # 构建系统
|
||||
├── _run.sh # 启动脚本
|
||||
├── _set.sh # 编译 + 链接 configs/
|
||||
└── README.md # 模块文档
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 6. 关键依赖
|
||||
|
||||
| 依赖 | 用途 | 已安装 |
|
||||
|------|------|--------|
|
||||
| `libmsg` (`-lus_msg`) | RpcServer / RpcClient / MsgServer / AsyncQueue | ✅ |
|
||||
| `boost::json` (`-lboost_json`) | JSON 序列化/反序列化 | ✅ |
|
||||
| `boost::program_options` (`-lboost_program_options`) | 命令行参数解析 | ✅ |
|
||||
| `pthread` | 多线程(mutex) | ✅ |
|
||||
| `std::filesystem` | 文件系统操作(C++20 标准库) | ✅ |
|
||||
|
||||
**零新增系统依赖**——所有库均已被 captcha2 / smtp3 等模块使用。
|
||||
|
||||
---
|
||||
|
||||
## 7. 测试计划
|
||||
|
||||
### 7.1 功能测试(对比 Python 版响应一致性)
|
||||
|
||||
| 测试项 | 方法 | 通过标准 |
|
||||
|--------|------|----------|
|
||||
| `get` 取全文 | `RpcClient.call("get", {"name":"db.ini"})` | content/size/mtime 与 Python 版一致 |
|
||||
| `get` INI 按键 | `RpcClient.call("get", {"name":"db.ini","key":"server.port"})` | value="8005" |
|
||||
| `get` JSON 按键 | 构造 .json 配置,取 `server.port` | 值一致 |
|
||||
| `get` 不存在 | `RpcClient.call("get", {"name":"nonexist.ini"})` | `success=false`,error 含"不存在" |
|
||||
| `get` 路径遍历 | `{"name":"../../etc/passwd"}` | `success=false`,拒绝 |
|
||||
| `get` 非法扩展名 | `{"name":"test.exe"}` | `success=false`,拒绝 |
|
||||
| `list` | `RpcClient.call("list", {})` | 配置列表与 Python 版一致 |
|
||||
| `reload` 指定 | `RpcClient.call("reload", {"name":"db.ini"})` | `success=true`,reloaded=["db.ini"] |
|
||||
| `reload` 全部 | `RpcClient.call("reload", {})` | `success=true`,reloaded 含全部缓存名 |
|
||||
| `reload` 空字符串 | `RpcClient.call("reload", {"name":""})` | `success=false`,error 含"不能为空" |
|
||||
| `reload` 非字符串 | `RpcClient.call("reload", {"name":123})` | `success=false`,error 含"必须是字符串" |
|
||||
| `health` | `RpcClient.call("health", {})` | status="ok",service="confs" |
|
||||
| mtime 重载 | 修改配置文件内容,再 `get` | 返回新内容 |
|
||||
| 大文件拒绝 | 创建 > 1MiB 文件,`get` | `success=false`,error 含"过大" |
|
||||
| 符号链接逃逸 | configs/ 内放逃逸符号链接 | `list` 中不出现 |
|
||||
|
||||
### 7.2 消费方回归测试
|
||||
|
||||
| 服务 | 测试方法 |
|
||||
|------|----------|
|
||||
| static | 启动 confs2 + static,验证静态文件服务正常(从 confs2 获取 static.ini) |
|
||||
| dnfs | 启动 confs2 + dnfs,验证树状菜单 API 正常 |
|
||||
| db | 启动 confs2 + db,验证数据库服务正常(从 confs2 获取 db.ini) |
|
||||
| captcha | 启动 confs2 + captcha2,验证验证码生成正常(从 confs2 获取 captcha.ini) |
|
||||
| redirect80 | 启动 confs2 + redirect80,验证重定向服务正常 |
|
||||
|
||||
### 7.3 安全测试
|
||||
|
||||
| 测试项 | 方法 |
|
||||
|--------|------|
|
||||
| 路径遍历 | `{"name":"../../../etc/passwd"}` → 拒绝 |
|
||||
| 绝对路径 | `{"name":"/etc/passwd"}` → 拒绝 |
|
||||
| 非法扩展名 | `{"name":"malicious.sh"}` → 拒绝 |
|
||||
| 符号链接逃逸 | configs/ 内符号链接指向 /etc/ → list 跳过,get 拒绝 |
|
||||
| 大文件 OOM | 创建 2MiB 配置 → 拒绝 |
|
||||
| 日志注入 | 配置名含 `\n[FAKE]` → 日志中无伪造条目 |
|
||||
|
||||
---
|
||||
|
||||
## 8. 安全考量
|
||||
|
||||
### 8.1 安全机制对照表
|
||||
|
||||
| 机制 | Python 版状态 | C++ 版方案 |
|
||||
|------|---------------|-------------|
|
||||
| 扩展名白名单 | ✅ | ✅ `kAllowedExt` |
|
||||
| 路径遍历防护 | ✅ basename + realpath 前缀 | ✅ `filename()` + `canonical()` 前缀 |
|
||||
| TOCTOU 防护 | ✅ 锁内 realpath→stat→open | ✅ 锁内 canonical→file_size→ifstream |
|
||||
| 文件大小限制 | ✅ 1MiB | ✅ `kMaxConfigSize = 1MiB` |
|
||||
| Slowloris 防护 | ✅ pymsg 10s recv 超时 | ✅ libmsg MsgServer 同策略 |
|
||||
| 缓存防污染 | ✅ 返回浅拷贝 | ✅ 返回值拷贝 |
|
||||
| list 符号链接防护 | ✅ realpath 前缀校验 | ✅ canonical 前缀校验 |
|
||||
| 参数类型校验 | ✅ 非字符串/空值报错 | ✅ 逐字段 is_string 校验 |
|
||||
| INI 安全 | ✅ 禁用插值 | ✅ 不实现插值 |
|
||||
| 仅监听本机 | ✅ 127.0.0.1 | ✅ 默认 127.0.0.1 |
|
||||
| 异常隔离 | ✅ RpcServer 包装 | ✅ libmsg RpcServer 同机制 |
|
||||
| 日志注入防护 | ❌ P2-3 遗留 | ✅ **新增** sanitize_for_log |
|
||||
| 异常消息脱敏 | ❌ P3-6 遗留 | ✅ **新增** RpcServer 内置脱敏 + 业务层包装 |
|
||||
| 鉴权 | ❌ 无 | ❌ 无(依赖仅本机信任边界,与 Python 版一致) |
|
||||
|
||||
### 8.2 INI 大小写策略文档化(修复 P3-5)
|
||||
|
||||
在代码注释和 README 中明确文档化:
|
||||
|
||||
```
|
||||
INI 解析大小写策略(与 Python configparser(interpolation=None) 一致):
|
||||
- Section 名:大小写敏感([Server] ≠ [server])
|
||||
- Option 名:大小写不敏感(port = Port = PORT)
|
||||
- 值:原样返回,不做任何转换
|
||||
- 插值:禁用,%(key)s 原样返回
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 9. 兼容性与迁移
|
||||
|
||||
### 9.1 零改动兼容
|
||||
|
||||
- **消费方零改动**:5 个 Python 服务继续用 `RpcClient("127.0.0.1", 8010, ...)` 调用,协议/端口/路由/响应格式全部不变。
|
||||
- **配置文件零迁移**:`confs2/configs/` 符号链接到 `../confs/configs/`,现有 9 个 .ini 文件不动。
|
||||
- **启动方式一致**:`init.txt` 中 `bash .../confs/_run.sh` 改为 `bash .../confs2/_run.sh`,一行改动。
|
||||
|
||||
### 9.2 迁移路径
|
||||
|
||||
```
|
||||
阶段 1: 编译 confs2 + 功能测试(confs2 与 confs 并存)
|
||||
↓
|
||||
阶段 2: 消费方回归测试(逐个服务切到 confs2 验证)
|
||||
↓
|
||||
阶段 3: init.txt 切换 confs → confs2(一行改动)
|
||||
↓
|
||||
阶段 4: 生产运行观察 72h,确认无异常
|
||||
↓
|
||||
阶段 5: Python confs 目录归档(不立即删除,保留回退能力)
|
||||
```
|
||||
|
||||
### 9.3 init.txt 切换
|
||||
|
||||
```
|
||||
# 切换前
|
||||
bash /opt/micro_servers/confs/_run.sh
|
||||
|
||||
# 切换后
|
||||
bash /opt/micro_servers/confs2/_run.sh
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 10. 未来扩展:core 注册中心联动
|
||||
|
||||
当 core 注册中心([core_service_registry_dev_doc](core_service_registry_dev_doc.md))落地后,confs2 可作为注册中心的第一个自注册服务:
|
||||
|
||||
```cpp
|
||||
// main() 启动后,RpcServer.start() 之前
|
||||
register_to_core("confs", g_config.port, "pymsg", "health",
|
||||
"bash /opt/micro_servers/confs2/_run.sh");
|
||||
```
|
||||
|
||||
进一步地,当未来评估 confs2 合入 core 时(方案 A: HTTP 合入),因 confs2 已是 C++ 同构,迁移成本仅为:
|
||||
1. 将 ConfigStore 代码移入 core 进程
|
||||
2. HTTP :5200 增加 `/confs/*` 路由
|
||||
3. 5 个消费方从 pymsg 改 HTTP(~50 行改动)
|
||||
|
||||
详见前序分析中的"confs2 合入 core 注册中心"可行性评估。
|
||||
|
||||
---
|
||||
|
||||
## 11. 附录
|
||||
|
||||
### A. 与已有 C++ 重写的模式对照
|
||||
|
||||
| 维度 | captcha2 | smtp3 | **confs2** |
|
||||
|------|----------|-------|-----------|
|
||||
| 原版语言 | Python | Python | Python |
|
||||
| 重写语言 | C++20 | C++20 | C++20 |
|
||||
| 协议 | pymsg(libmsg) | pymsg(libmsg) | pymsg(libmsg) |
|
||||
| 端口 | 8009 | 8008 | 8010 |
|
||||
| libmsg | ✅ RpcServer | ✅ RpcServer | ✅ RpcServer |
|
||||
| boost_json | ✅ | ✅ | ✅ |
|
||||
| LogC.hpp | ✅ 原创 | ✅ 改编 | ✅ 改编 |
|
||||
| Makefile 模式 | g++ -std=c++20 | g++ -std=c++20 | g++ -std=c++20 |
|
||||
| _run.sh 模式 | sleep 1 + exec .bin | sleep 1 + exec .bin | sleep 1 + exec .bin |
|
||||
| _set.sh 模式 | set -euo + make | set -euo + make | set -euo + make + ln -s |
|
||||
| 特殊依赖 | freetype, stb | openssl | 无(仅 libmsg + boost) |
|
||||
|
||||
### B. 配置文件目录现状
|
||||
|
||||
| 文件名 | 消费方 | 大小 |
|
||||
|--------|--------|------|
|
||||
| ca.ini | ca 模块 | 51 行 |
|
||||
| captcha.ini | captcha2 | 11 行 |
|
||||
| db.ini | db | 16 行 |
|
||||
| dnfs.ini | dnfs | 16 行 |
|
||||
| example.ini | 示例 | 10 行 |
|
||||
| gateway.ini | gateway | 20 行 |
|
||||
| redirect80.ini | redirect80 | 5 行 |
|
||||
| smtp.ini | smtp3 | 19 行 |
|
||||
| static.ini | static | 17 行 |
|
||||
|
||||
### C. libmsg RpcServer API 速查
|
||||
|
||||
```cpp
|
||||
// 构造
|
||||
pmc::net::RpcServer rpc("127.0.0.1", 8010,
|
||||
pmc::net::MsgServerConfig{.workers = 4},
|
||||
[](const std::string& route, const std::exception& e) {
|
||||
// on_error 回调
|
||||
});
|
||||
|
||||
// 注册路由
|
||||
rpc.add_route("get", [](const json::object& req) -> json::value {
|
||||
// handler 签名: json::value(const json::object& req)
|
||||
// 返回值自动包装为 {"success": true, "data": ...}
|
||||
// 抛异常自动包装为 {"success": false, "error": "<脱敏>"}
|
||||
return json::object{{"content", "..."}};
|
||||
});
|
||||
|
||||
// 启动(阻塞)
|
||||
rpc.start();
|
||||
|
||||
// 停止(另一线程)
|
||||
rpc.stop();
|
||||
```
|
||||
@@ -1,806 +0,0 @@
|
||||
# Core 启动配置 YAML 化与 CLI 参数标准化 · 开发文档
|
||||
|
||||
> 文档类型:功能开发计划
|
||||
> 模块:core(PMC / Nidus)+ 全部子服务
|
||||
> 作者:自动生成
|
||||
> 日期:2026-08-10
|
||||
> 前置关系:本文档是 `core_service_registry_dev_doc.md` P7 阶段(services.yaml)的前置拆分,聚焦"配置格式迁移 + 参数透传 + CLI 标准化"三步基础设施改造
|
||||
|
||||
---
|
||||
|
||||
## 1. 概述
|
||||
|
||||
### 1.1 背景与动机
|
||||
|
||||
当前 core 通过 `init.txt`(纯命令列表)拉起 11 个子服务,存在以下问题:
|
||||
|
||||
| 现状 | 问题 |
|
||||
|------|------|
|
||||
| `init.txt` 每行一条 bash 命令,无结构化字段 | 无法为每个服务指定命令行参数、环境变量、依赖关系 |
|
||||
| 5/11 个 `_run.sh` 不转发 `"$@"` | core 即便想传参也传不进去 |
|
||||
| `sleep 2` 等依赖等待逻辑散落在各 `_run.sh` 中 | 无统一管理,调整需逐个脚本修改 |
|
||||
| waf 的 `-c config.yaml` 硬编码在 `_run.sh` | 参数不可外部覆盖 |
|
||||
| 各服务 CLI 接口不统一 | ca/gateway 仅有 `--help/--version`,waf 有 `-c/-list/-test/-check`,其余 8 个服务零 CLI 支持 |
|
||||
|
||||
**本方案目标**:通过三步改造,让 core 能够以结构化 YAML 声明每个服务的启动命令和参数,参数能透传到各服务二进制/脚本,并建立统一的 CLI 参数规范——为注册中心 `services.yaml` + `depends_on` DAG 拓扑排序打下基础。
|
||||
|
||||
### 1.2 设计目标
|
||||
|
||||
1. **结构化启动声明**:`init.txt` → `services.yaml`,每个服务可声明 `cmd`、`args`、`env`、`depends_on`、`sleep_before`
|
||||
2. **参数透传**:全部 11 个 `_run.sh` 统一转发 `"$@"`,消除硬编码参数
|
||||
3. **CLI 标准化**:定义通用参数约定(`--host`/`--port`/`--config`/`--log-level`/`--instance`),各服务逐步适配
|
||||
4. **向后兼容**:`services.yaml` 不存在时回退 `init.txt`,现有部署零中断
|
||||
5. **yaml-cpp 引入**:vendored 到 `core/third_party/yaml-cpp/`,不依赖系统包管理器
|
||||
|
||||
### 1.3 与注册中心计划的关系
|
||||
|
||||
```
|
||||
本文档(三步改造) 注册中心计划(core_service_registry_dev_doc.md)
|
||||
┌─────────────────────┐ ┌─────────────────────────────┐
|
||||
│ Step 1: services.yaml│ ──────▶ │ P7: services.yaml DAG 排序 │
|
||||
│ Step 2: _run.sh 透传 │ ──────▶ │ P2: init 解析后自动填充注册表 │
|
||||
│ Step 3: CLI 标准化 │ ──────▶ │ P5: RegistrationClient SDK │
|
||||
└─────────────────────┘ └─────────────────────────────┘
|
||||
```
|
||||
|
||||
本方案是注册中心的**基础设施前置**——做完这三步,P2/P5/P7 的工作量直接减半。
|
||||
|
||||
---
|
||||
|
||||
## 2. 现状分析
|
||||
|
||||
### 2.1 init.txt 当前内容
|
||||
|
||||
```
|
||||
bash /opt/micro_servers/db2/_run.sh
|
||||
bash /opt/micro_servers/ai_review/_run.sh
|
||||
bash /opt/micro_servers/dnfs2/_run.sh
|
||||
bash /opt/micro_servers/smtp/_run.sh
|
||||
bash /opt/micro_servers/fsys/_run.sh
|
||||
bash /opt/micro_servers/waf/_run.sh
|
||||
bash /opt/micro_servers/ca/_run.sh
|
||||
bash /opt/micro_servers/gateway/_run.sh
|
||||
bash /opt/micro_servers/captcha2/_run.sh
|
||||
bash /opt/micro_servers/aibox/_run.sh
|
||||
bash /opt/micro_servers/redirect80/_run.sh
|
||||
```
|
||||
|
||||
纯命令列表,无参数、无元数据、无依赖声明。
|
||||
|
||||
### 2.2 各服务 _run.sh 现状
|
||||
|
||||
| 服务 | 语言 | `"$@"` 转发 | 硬编码参数 | 特殊处理 | sleep 等待 |
|
||||
|------|------|-------------|-----------|----------|------------|
|
||||
| db2 | C++ | ✅ `"$@"` | — | `set -e` | — |
|
||||
| dnfs2 | C++ | ✅ `"$@"` | — | `set -e` | — |
|
||||
| fsys | C++ | ✅ `"$@"` | — | `set -e` | — |
|
||||
| ca | C++ | ✅ `"$@"` | — | — | `sleep 2` |
|
||||
| gateway | C++ | ✅ `"$@"` | — | — | `sleep 2` |
|
||||
| waf | Go | ❌ | `-c config.yaml` | — | — |
|
||||
| captcha2 | C++ | ❌ | — | — | `sleep 1` |
|
||||
| aibox | Python | ⚠️ `$@`(未加引号) | — | `_setenv.sh` + venv | — |
|
||||
| ai_review | Python | ❌ | — | venv | — |
|
||||
| smtp | Python | ❌ | — | venv | `sleep 2` |
|
||||
| redirect80 | Python | ❌ | — | setcap + venv | — |
|
||||
|
||||
**统计**:5 个已转发、1 个半转发(未加引号)、5 个未转发。3 个脚本含 `sleep` 硬编码。
|
||||
|
||||
### 2.3 各服务 CLI 参数支持现状
|
||||
|
||||
| 服务 | 语言 | CLI 框架 | 实际参数 | 配置来源 |
|
||||
|------|------|----------|----------|----------|
|
||||
| db2 | C++ | 无 | 无 | 硬编码常量 |
|
||||
| dnfs2 | C++ | 无 | 无 | 硬编码 + confs |
|
||||
| fsys | C++ | 无 | 无 | 硬编码 |
|
||||
| captcha2 | C++ | 无 | 无 | 硬编码 |
|
||||
| ca | C++ | `boost::program_options` | `--help` `--version`(仅此) | confs (ca.ini) + 环境变量 |
|
||||
| gateway | C++ | `boost::program_options` | `--help` `--version`(仅此) | confs (gateway.ini) + 环境变量 |
|
||||
| waf | Go | `flag` | `-c` `-list` `-test` `-check` | YAML 配置文件 |
|
||||
| aibox | Python | `argparse`(已 import 未使用) | 无 | `_setenv.sh` + 硬编码 |
|
||||
| ai_review | Python | 无 | 无 | 硬编码 |
|
||||
| smtp | Python | 无 | 无 | 硬编码 |
|
||||
| redirect80 | Python | 无 | 无 | 硬编码 |
|
||||
|
||||
**关键发现**:
|
||||
- ca 和 gateway 已引入 `boost::program_options` 但只实现了 `--help`/`--version`,实际配置全走 confs + 环境变量
|
||||
- aibox 已 `import argparse` 但从未调用,是死代码
|
||||
- waf 是唯一有完整 CLI 接口的服务(Go `flag` 包)
|
||||
- 8/11 个服务零 CLI 参数支持
|
||||
|
||||
### 2.4 配置加载模式分类
|
||||
|
||||
当前各服务的配置获取有三种模式:
|
||||
|
||||
| 模式 | 服务 | 说明 |
|
||||
|------|------|------|
|
||||
| **confs + 环境变量** | ca, gateway, dnfs2 | 从 confs 服务拉取 ini 配置,环境变量可覆盖;`CONFS_HOST`/`CONFS_PORT` 引导 |
|
||||
| **本地配置文件** | waf | `-c config.yaml` 指定 YAML 配置 |
|
||||
| **硬编码** | db2, fsys, captcha2, aibox, ai_review, smtp, redirect80 | 常量直接写在源码中 |
|
||||
|
||||
CLI 标准化后,三种模式可统一为:**CLI 参数(最高优先级)> 环境变量 > confs/配置文件 > 硬编码默认值**。
|
||||
|
||||
---
|
||||
|
||||
## 3. 三步改造方案
|
||||
|
||||
### 3.1 总览
|
||||
|
||||
```
|
||||
Step 1: init.txt → services.yaml Step 2: _run.sh 参数透传 Step 3: CLI 参数标准化
|
||||
┌──────────────────────────────┐ ┌──────────────────────┐ ┌──────────────────────┐
|
||||
│ • 引入 yaml-cpp (vendored) │ │ • 11 个 _run.sh 统一 │ │ • 定义通用参数约定 │
|
||||
│ • 设计 services.yaml 格式 │ ──▶ │ 末尾 exec 加 "$@" │ ──▶ │ • _run.sh 注入默认值 │
|
||||
│ • core 解析 YAML 并启动 │ │ • 移除硬编码参数 │ │ • 各服务逐步适配 │
|
||||
│ • 回退兼容 init.txt │ │ • sleep 移至 YAML │ │ • 环境变量桥接 │
|
||||
└──────────────────────────────┘ └──────────────────────┘ └──────────────────────┘
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 4. Step 1:init.txt → services.yaml
|
||||
|
||||
### 4.1 yaml-cpp 引入
|
||||
|
||||
采用系统包安装方式(已写入 `_setup.deb.sh` 和 `_setup.rhel.sh`):
|
||||
|
||||
| 发行版 | 包名 | 安装命令 |
|
||||
|--------|------|----------|
|
||||
| Debian/Ubuntu | `libyaml-cpp-dev` | `apt install libyaml-cpp-dev` |
|
||||
| RHEL/CentOS/Rocky | `yaml-cpp-devel` | `dnf install yaml-cpp-devel` |
|
||||
|
||||
**Makefile 集成**(已完成):
|
||||
|
||||
```makefile
|
||||
# 链接 yaml-cpp 系统库
|
||||
LIBS := ... -lyaml-cpp
|
||||
|
||||
# InitParser 编译目标
|
||||
$(BUILD_DIR)/InitParser.o: src/registry/InitParser.cpp
|
||||
$(CPP) $(CPPFLAGS) -c $^ $(INCS) -o $@
|
||||
|
||||
# 链接时追加
|
||||
$(APP_NAME).out: ... $(BUILD_DIR)/InitParser.o
|
||||
```
|
||||
|
||||
### 4.2 services.yaml 格式设计
|
||||
|
||||
```yaml
|
||||
# core/services.yaml
|
||||
# 声明式服务启动配置(替代 init.txt)
|
||||
# 如果本文件存在则优先使用,否则回退 init.txt
|
||||
|
||||
services:
|
||||
# --- 数据层 ---
|
||||
- name: db2
|
||||
cmd: bash /opt/micro_servers/db2/_run.sh
|
||||
args: []
|
||||
env: {}
|
||||
depends_on: []
|
||||
sleep_before: 0
|
||||
|
||||
# --- AI 服务 ---
|
||||
- name: ai_review
|
||||
cmd: bash /opt/micro_servers/ai_review/_run.sh
|
||||
args: ["--port", "8007", "--log-level", "info"]
|
||||
env:
|
||||
PYTHONUNBUFFERED: "1"
|
||||
depends_on: [db2]
|
||||
sleep_before: 2
|
||||
|
||||
# --- 文件服务 ---
|
||||
- name: dnfs2
|
||||
cmd: bash /opt/micro_servers/dnfs2/_run.sh
|
||||
args: []
|
||||
env: {}
|
||||
depends_on: [db2, ai_review]
|
||||
sleep_before: 0
|
||||
|
||||
# --- 邮件服务 ---
|
||||
- name: smtp
|
||||
cmd: bash /opt/micro_servers/smtp/_run.sh
|
||||
args: []
|
||||
env: {}
|
||||
depends_on: []
|
||||
sleep_before: 2
|
||||
|
||||
# --- 前端静态资源 ---
|
||||
- name: fsys
|
||||
cmd: bash /opt/micro_servers/fsys/_run.sh
|
||||
args: []
|
||||
env: {}
|
||||
depends_on: []
|
||||
sleep_before: 0
|
||||
|
||||
# --- 安全链路 ---
|
||||
- name: waf
|
||||
cmd: bash /opt/micro_servers/waf/_run.sh
|
||||
args: ["-c", "config.yaml"]
|
||||
env: {}
|
||||
depends_on: [ca, gateway]
|
||||
sleep_before: 0
|
||||
|
||||
- name: ca
|
||||
cmd: bash /opt/micro_servers/ca/_run.sh
|
||||
args: []
|
||||
env: {}
|
||||
depends_on: [db2]
|
||||
sleep_before: 2
|
||||
|
||||
- name: gateway
|
||||
cmd: bash /opt/micro_servers/gateway/_run.sh
|
||||
args: []
|
||||
env: {}
|
||||
depends_on: [db2]
|
||||
sleep_before: 2
|
||||
|
||||
# --- 验证码 ---
|
||||
- name: captcha2
|
||||
cmd: bash /opt/micro_servers/captcha2/_run.sh
|
||||
args: []
|
||||
env: {}
|
||||
depends_on: []
|
||||
sleep_before: 1
|
||||
|
||||
# --- AI 工具箱 ---
|
||||
- name: aibox
|
||||
cmd: bash /opt/micro_servers/aibox/_run.sh
|
||||
args: []
|
||||
env:
|
||||
PATH: "${PATH}:~/.aibox/skills"
|
||||
depends_on: []
|
||||
sleep_before: 0
|
||||
|
||||
# --- 端口重定向 ---
|
||||
- name: redirect80
|
||||
cmd: bash /opt/micro_servers/redirect80/_run.sh
|
||||
args: []
|
||||
env: {}
|
||||
depends_on: [waf]
|
||||
sleep_before: 0
|
||||
```
|
||||
|
||||
**字段说明**:
|
||||
|
||||
| 字段 | 类型 | 必填 | 说明 |
|
||||
|------|------|------|------|
|
||||
| `name` | string | 是 | 服务唯一标识,用于日志和未来注册表 |
|
||||
| `cmd` | string | 是 | 启动命令(与 init.txt 每行格式一致) |
|
||||
| `args` | string[] | 否 | 传递给 `_run.sh` 的命令行参数列表 |
|
||||
| `env` | map<string,string> | 否 | 启动前设置的环境变量 |
|
||||
| `depends_on` | string[] | 否 | 依赖的服务名列表(未来 DAG 排序用,当前仅记录) |
|
||||
| `sleep_before` | int | 否 | 启动前等待秒数(替代散落在 _run.sh 中的 sleep) |
|
||||
|
||||
### 4.3 core 解析逻辑
|
||||
|
||||
新增 `include/registry/InitParser.hh` + `src/registry/InitParser.cpp`:
|
||||
|
||||
```cpp
|
||||
// include/registry/InitParser.hh
|
||||
#pragma once
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <map>
|
||||
|
||||
namespace qing::registry {
|
||||
|
||||
struct ServiceConfig {
|
||||
std::string name;
|
||||
std::string cmd;
|
||||
std::vector<std::string> args;
|
||||
std::map<std::string, std::string> env;
|
||||
std::vector<std::string> depends_on;
|
||||
int sleep_before = 0;
|
||||
};
|
||||
|
||||
class InitParser {
|
||||
public:
|
||||
// 优先读 services.yaml,不存在则回退 init.txt
|
||||
static std::vector<ServiceConfig> load(const std::string& conf_dir);
|
||||
|
||||
// 解析 services.yaml
|
||||
static std::vector<ServiceConfig> parseYaml(const std::string& path);
|
||||
|
||||
// 解析 init.txt(向后兼容)
|
||||
static std::vector<ServiceConfig> parseTxt(const std::string& path);
|
||||
|
||||
// 将 ServiceConfig 拼成完整命令字符串(cmd + " " + args.join(" "))
|
||||
static std::string toCommand(const ServiceConfig& svc);
|
||||
};
|
||||
|
||||
} // namespace qing::registry
|
||||
```
|
||||
|
||||
**解析流程**:
|
||||
|
||||
```
|
||||
load(conf_dir):
|
||||
1. 检查 conf_dir/services.yaml 是否存在
|
||||
2. 存在 → parseYaml() → 返回 ServiceConfig 列表
|
||||
3. 不存在 → 检查 conf_dir/init.txt
|
||||
4. 存在 → parseTxt() → 每行作为 ServiceConfig{name="", cmd=line, args=[]}
|
||||
5. 都不存在 → 报错退出
|
||||
```
|
||||
|
||||
**parseYaml 核心逻辑**(使用 yaml-cpp):
|
||||
|
||||
```cpp
|
||||
#include <yaml-cpp/yaml.h>
|
||||
|
||||
std::vector<ServiceConfig> InitParser::parseYaml(const std::string& path) {
|
||||
YAML::Node root = YAML::LoadFile(path);
|
||||
std::vector<ServiceConfig> result;
|
||||
|
||||
for (const auto& node : root["services"]) {
|
||||
ServiceConfig svc;
|
||||
svc.name = node["name"].as<std::string>("");
|
||||
svc.cmd = node["cmd"].as<std::string>();
|
||||
|
||||
if (node["args"])
|
||||
for (const auto& a : node["args"])
|
||||
svc.args.push_back(a.as<std::string>());
|
||||
|
||||
if (node["env"])
|
||||
for (const auto& kv : node["env"])
|
||||
svc.env[kv.first.as<std::string>()] = kv.second.as<std::string>();
|
||||
|
||||
if (node["depends_on"])
|
||||
for (const auto& d : node["depends_on"])
|
||||
svc.depends_on.push_back(d.as<std::string>());
|
||||
|
||||
if (node["sleep_before"])
|
||||
svc.sleep_before = node["sleep_before"].as<int>(0);
|
||||
|
||||
result.push_back(std::move(svc));
|
||||
}
|
||||
return result;
|
||||
}
|
||||
```
|
||||
|
||||
**启动执行逻辑**(改动 `pmc_mtd::parse_self_init_list`):
|
||||
|
||||
```cpp
|
||||
// pmc_mtd.cpp 中改造
|
||||
void pmc_mtd::init_from_config(const std::string& conf_dir) {
|
||||
auto services = InitParser::load(conf_dir);
|
||||
|
||||
for (const auto& svc : services) {
|
||||
// 1. sleep_before 等待
|
||||
if (svc.sleep_before > 0)
|
||||
std::this_thread::sleep_for(std::chrono::seconds(svc.sleep_before));
|
||||
|
||||
// 2. 设置环境变量
|
||||
for (const auto& [k, v] : svc.env)
|
||||
setenv(k.c_str(), v.c_str(), 1);
|
||||
|
||||
// 3. 拼接完整命令并启动
|
||||
std::string full_cmd = InitParser::toCommand(svc);
|
||||
// full_cmd = "bash /opt/micro_servers/db2/_run.sh --port 8007 --log-level info"
|
||||
ppool.crtp(full_cmd);
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### 4.4 向后兼容策略
|
||||
|
||||
```
|
||||
services.yaml 存在?
|
||||
│
|
||||
┌─── 是 ──┤── 否 ───┐
|
||||
▼ ▼
|
||||
parseYaml() init.txt 存在?
|
||||
│ │
|
||||
│ ┌── 是 ──┤── 否 ───┐
|
||||
│ ▼ ▼
|
||||
│ parseTxt() 报错退出
|
||||
│ │
|
||||
└─────┬─────┘
|
||||
▼
|
||||
ServiceConfig[]
|
||||
逐个启动
|
||||
```
|
||||
|
||||
- `init.txt` 模式下,每行解析为 `{name="", cmd=line, args=[], sleep_before=0}`,行为与当前完全一致
|
||||
- `services.yaml` 模式下,享受参数透传和 sleep 管理能力
|
||||
- Makefile 的 `CONF_NAME` 可逐步从 `init.txt` 切换为 `services.yaml`
|
||||
|
||||
---
|
||||
|
||||
## 5. Step 2:_run.sh 参数透传
|
||||
|
||||
### 5.1 改造规则
|
||||
|
||||
每个 `_run.sh` 的末尾 `exec` 行必须带 `"$@"`,同时移除硬编码参数和散落的 `sleep`。
|
||||
|
||||
**改造模板**:
|
||||
|
||||
```bash
|
||||
#!/bin/bash
|
||||
# <service>/_run.sh — 标准化启动脚本
|
||||
cd "$(dirname "$0")" || exit
|
||||
|
||||
# --- 服务特有前置(venv / setcap 等)保持不变 ---
|
||||
|
||||
# 启动:"$@" 透传 core/services.yaml 传入的参数
|
||||
exec <binary_or_python> "$@"
|
||||
```
|
||||
|
||||
### 5.2 逐服务改造清单
|
||||
|
||||
| 服务 | 当前 exec 行 | 改造后 | 额外改动 |
|
||||
|------|-------------|--------|----------|
|
||||
| db2 | `exec ./build/db2 "$@"` | ✅ 无需改 | — |
|
||||
| dnfs2 | `exec ./build/dnfs2 "$@"` | ✅ 无需改 | — |
|
||||
| fsys | `exec ./build/fsys "$@"` | ✅ 无需改 | — |
|
||||
| ca | `exec ./build/ca.out "$@"` | ✅ 无需改 | 移除 `sleep 2`(移交 YAML) |
|
||||
| gateway | `exec ./build/gateway.out "$@"` | ✅ 无需改 | 移除 `sleep 2`(移交 YAML) |
|
||||
| waf | `exec ./waf.bin -c config.yaml` | `exec ./waf.bin -c config.yaml "$@"` | 保留 `-c config.yaml` 作为脚本默认值,services.yaml `args: []` |
|
||||
| captcha2 | `exec ./captcha2.bin` | `exec ./captcha2.bin "$@"` | 移除 `sleep 1`(移交 YAML) |
|
||||
| aibox | `exec python aibox.py $@` | `exec python aibox.py "$@"` | `$@` → `"$@"` 加引号 |
|
||||
| ai_review | `exec python3 ai_review.py` | `exec python3 ai_review.py "$@"` | — |
|
||||
| smtp | `exec python3 smtp_v2.py` | `exec python3 smtp_v2.py "$@"` | 移除 `sleep 2`(移交 YAML) |
|
||||
| redirect80 | `exec python3 redirect80.py` | `exec python3 redirect80.py "$@"` | — |
|
||||
|
||||
**统计**:6 个无需改、5 个需改 exec 行、4 个需移除 sleep。
|
||||
|
||||
### 5.3 sleep 迁移对照
|
||||
|
||||
| 服务 | 当前 sleep 位置 | 迁移到 services.yaml |
|
||||
|------|----------------|---------------------|
|
||||
| ca | `_run.sh` 中 `sleep 2` | `sleep_before: 2` |
|
||||
| gateway | `_run.sh` 中 `sleep 2` | `sleep_before: 2` |
|
||||
| captcha2 | `_run.sh` 中 `sleep 1` | `sleep_before: 1` |
|
||||
| smtp | `_run.sh` 中 `sleep 2` | `sleep_before: 2` |
|
||||
|
||||
移除后 `_run.sh` 只负责"切换目录 + 前置准备 + exec",不负责时序控制。
|
||||
|
||||
---
|
||||
|
||||
## 6. Step 3:CLI 参数标准化
|
||||
|
||||
### 6.1 通用参数约定
|
||||
|
||||
定义一组所有服务都应支持的通用 CLI 参数:
|
||||
|
||||
| 参数 | 短选项 | 类型 | 默认值 | 说明 |
|
||||
|------|--------|------|--------|------|
|
||||
| `--host` | `-H` | string | `127.0.0.1` | 绑定地址 |
|
||||
| `--port` | `-p` | int | 服务特定 | 监听端口 |
|
||||
| `--config` | `-c` | string | 服务特定 | 配置文件路径 |
|
||||
| `--log-level` | `-l` | enum | `info` | `debug\|info\|warn\|error` |
|
||||
| `--instance` | `-i` | string | 服务名 | 实例标识(未来多实例/注册中心用) |
|
||||
| `--help` | `-h` | flag | — | 输出帮助 |
|
||||
| `--version` | `-v` | flag | — | 输出版本 |
|
||||
|
||||
**优先级规则**:`CLI 参数 > 环境变量 > confs/配置文件 > 硬编码默认值`
|
||||
|
||||
### 6.2 _run.sh 默认值注入
|
||||
|
||||
在服务二进制还不支持所有参数的过渡期,`_run.sh` 可注入默认值:
|
||||
|
||||
```bash
|
||||
#!/bin/bash
|
||||
# waf/_run.sh — 标准化后
|
||||
cd "$(dirname "$0")" || exit
|
||||
|
||||
# 默认参数(可被 services.yaml args 覆盖)
|
||||
DEFAULT_ARGS=(-c config.yaml)
|
||||
|
||||
exec ./waf.bin "${DEFAULT_ARGS[@]}" "$@"
|
||||
```
|
||||
|
||||
```bash
|
||||
#!/bin/bash
|
||||
# db2/_run.sh — 标准化后
|
||||
cd "$(dirname "$0")" || exit
|
||||
set -e
|
||||
|
||||
# 默认参数
|
||||
DEFAULT_ARGS=(--host 127.0.0.1 --port 8005 --log-level info)
|
||||
|
||||
exec ./build/db2 "${DEFAULT_ARGS[@]}" "$@"
|
||||
```
|
||||
|
||||
**原理**:`"$@"` 追加在 `DEFAULT_ARGS` 之后,如果用户通过 services.yaml 传了同名参数,取决于各服务的解析逻辑(`boost::program_options` 默认后者覆盖前者;Python `argparse` 同理)。
|
||||
|
||||
### 6.3 各语言适配指南
|
||||
|
||||
#### C++ 服务(db2, dnfs2, fsys, captcha2)
|
||||
|
||||
当前这 4 个服务无 CLI 框架。建议引入 `boost::program_options`(ca/gateway 已有先例):
|
||||
|
||||
```cpp
|
||||
// db2/src/db2.cpp — 新增 CLI 解析
|
||||
#include <boost/program_options.hpp>
|
||||
namespace po = boost::program_options;
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
po::options_description desc("db2 — database RPC service");
|
||||
desc.add_options()
|
||||
("help,h", "output help message")
|
||||
("version,v", "output version message")
|
||||
("host,H", po::value<std::string>()->default_value("127.0.0.1"), "bind address")
|
||||
("port,p", po::value<int>()->default_value(8005), "listen port")
|
||||
("log-level,l", po::value<std::string>()->default_value("info"),
|
||||
"debug|info|warn|error")
|
||||
("instance,i", po::value<std::string>()->default_value("db2"), "instance name")
|
||||
;
|
||||
|
||||
po::variables_map vm;
|
||||
try {
|
||||
po::store(po::parse_command_line(argc, argv, desc), vm);
|
||||
po::notify(vm);
|
||||
} catch (const std::exception& e) {
|
||||
std::cerr << "ERROR: " << e.what() << std::endl;
|
||||
std::cerr << desc << std::endl;
|
||||
return 1;
|
||||
}
|
||||
|
||||
if (vm.count("help")) { std::cout << desc << std::endl; return 0; }
|
||||
if (vm.count("version")) { std::cout << "db2 1.0.0" << std::endl; return 0; }
|
||||
|
||||
std::string host = vm["host"].as<std::string>();
|
||||
int port = vm["port"].as<int>();
|
||||
std::string log_level = vm["log-level"].as<std::string>();
|
||||
|
||||
// 后续使用 host/port/log_level 启动服务...
|
||||
}
|
||||
```
|
||||
|
||||
#### C++ 服务(ca, gateway)— 扩展现有 program_options
|
||||
|
||||
ca 和 gateway已有 `po::options_description` 但只注册了 `--help`/`--version`,需追加通用参数:
|
||||
|
||||
```cpp
|
||||
// ca/src/main.cpp — 扩展
|
||||
desc.add_options()
|
||||
("help,h", "output help message")
|
||||
("version,v", "output version message")
|
||||
// 新增通用参数
|
||||
("host,H", po::value<std::string>()->default_value("127.0.0.1"), "bind address")
|
||||
("port,p", po::value<int>()->default_value(8443), "listen port")
|
||||
("log-level,l", po::value<std::string>()->default_value("info"),
|
||||
"debug|info|warn|error")
|
||||
("instance,i", po::value<std::string>()->default_value("ca"), "instance name")
|
||||
;
|
||||
```
|
||||
|
||||
注意:ca/gateway 当前配置走 confs + 环境变量,CLI 参数应作为**覆盖层**,不替换 confs 机制。
|
||||
|
||||
#### Go 服务(waf)— 扩展现有 flag
|
||||
|
||||
waf 已有 `flag` 包,追加通用参数:
|
||||
|
||||
```go
|
||||
func main() {
|
||||
var (
|
||||
configFile = flag.String("c", "config.yaml", "配置文件路径")
|
||||
host = flag.String("H", "127.0.0.1", "绑定地址")
|
||||
port = flag.Int("p", 443, "监听端口")
|
||||
logLevel = flag.String("l", "info", "debug|info|warn|error")
|
||||
instance = flag.String("i", "waf", "实例名")
|
||||
doList = flag.Bool("list", false, "列出所有规则后退出")
|
||||
doTest = flag.String("test", "", "校验单个规则文件语法后退出")
|
||||
doCheck = flag.String("check", "", "用当前规则测试输入字符串后退出")
|
||||
)
|
||||
flag.Parse()
|
||||
// ...
|
||||
}
|
||||
```
|
||||
|
||||
注意:Go `flag` 包不支持短选项连写(`-H 127.0.0.1`),与 C++ `program_options` 的 `-H` 语法兼容。
|
||||
|
||||
#### Python 服务(aibox, ai_review, smtp, redirect80)
|
||||
|
||||
建议统一使用 `argparse`:
|
||||
|
||||
```python
|
||||
#!/usr/bin/env python3
|
||||
"""ai_review — AI 审核服务"""
|
||||
import argparse
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="ai_review — AI 审核服务")
|
||||
parser.add_argument("-H", "--host", default="127.0.0.1", help="绑定地址")
|
||||
parser.add_argument("-p", "--port", type=int, default=8007, help="监听端口")
|
||||
parser.add_argument("-l", "--log-level", default="info",
|
||||
choices=["debug", "info", "warn", "error"], help="日志级别")
|
||||
parser.add_argument("-i", "--instance", default="ai_review", help="实例名")
|
||||
return parser.parse_args()
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
# 后续使用 args.host, args.port, args.log_level...
|
||||
```
|
||||
|
||||
aibox 已有 `import argparse`(死代码),激活即可。
|
||||
|
||||
### 6.4 环境变量桥接
|
||||
|
||||
对于走 confs + 环境变量的服务(ca, gateway, dnfs2),CLI 参数应映射到环境变量,保持与现有 confs 机制兼容:
|
||||
|
||||
```bash
|
||||
# ca/_run.sh — 环境变量桥接模式
|
||||
cd "$(dirname "$0")" || exit
|
||||
|
||||
# 从命令行参数提取并导出环境变量(供 confs 加载逻辑使用)
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case "$1" in
|
||||
-H|--host) export CA_HOST="$2"; shift 2 ;;
|
||||
-p|--port) export CA_PORT="$2"; shift 2 ;;
|
||||
-l|--log-level) export LOG_LEVEL="$2"; shift 2 ;;
|
||||
*) EXTRA_ARGS+=("$1"); shift ;;
|
||||
esac
|
||||
done
|
||||
|
||||
exec ./build/ca.out "${EXTRA_ARGS[@]}"
|
||||
```
|
||||
|
||||
这种桥接方式不需要立刻改 ca 的 C++ 源码——`_run.sh` 先把 CLI 参数翻译成环境变量,ca 现有的环境变量读取逻辑自然生效。
|
||||
|
||||
---
|
||||
|
||||
## 7. 实施计划
|
||||
|
||||
### 7.1 分步交付
|
||||
|
||||
| 步骤 | 改动范围 | 涉及文件 | 预估工时 | 风险 |
|
||||
|------|----------|----------|----------|------|
|
||||
| **S1-a** | yaml-cpp vendored + Makefile 集成 | `core/third_party/yaml-cpp/`、`core/Makefile` | 1h | 低 |
|
||||
| **S1-b** | InitParser 实现 | `core/include/registry/InitParser.hh`、`core/src/registry/InitParser.cpp` | 2h | 低 |
|
||||
| **S1-c** | pmc_mtd 改造(调用 InitParser) | `core/src/main/pmc_mtd.cpp` | 1h | 中(改核心启动路径) |
|
||||
| **S1-d** | 创建 services.yaml | `core/services.yaml` | 0.5h | 低 |
|
||||
| **S1-e** | 回退测试(删除 services.yaml 验证 init.txt 兼容) | — | 0.5h | 低 |
|
||||
| **S2-a** | 5 个 _run.sh 改 exec 行 | `waf/`、`captcha2/`、`aibox/`、`ai_review/`、`smtp/`、`redirect80/` 的 `_run.sh` | 0.5h | 极低 |
|
||||
| **S2-b** | 4 个 _run.sh 移除 sleep | `ca/`、`gateway/`、`captcha2/`、`smtp/` 的 `_run.sh` | 0.2h | 低(sleep 已移交 YAML) |
|
||||
| **S3-a** | 定义 CLI 标准文档(本文档 §6) | — | 0h(已完成) | — |
|
||||
| **S3-b** | C++ 服务适配(ca/gateway 扩展,db2/dnfs2/fsys/captcha2 新增) | 各服务 `src/*.cpp` | 4h | 中 |
|
||||
| **S3-c** | Python 服务适配(argparse) | `aibox.py`、`ai_review.py`、`smtp_v2.py`、`redirect80.py` | 2h | 低 |
|
||||
| **S3-d** | Go 服务适配(waf flag 扩展) | `waf/waf.go` | 0.5h | 低 |
|
||||
| **S3-e** | _run.sh 默认值注入 | 各服务 `_run.sh` | 0.5h | 低 |
|
||||
| **S4** | 端到端测试 | — | 1h | — |
|
||||
|
||||
**总计预估:~14 小时**
|
||||
|
||||
### 7.2 推荐实施顺序
|
||||
|
||||
```
|
||||
S1-a (yaml-cpp 集成)
|
||||
└─▶ S1-b (InitParser)
|
||||
└─▶ S1-c (pmc_mtd 改造)
|
||||
└─▶ S1-d (services.yaml)
|
||||
└─▶ S1-e (回退测试) ──▶ Step 1 完成
|
||||
│
|
||||
▼
|
||||
S2-a + S2-b ( _run.sh 改造) ──▶ Step 2 完成
|
||||
│
|
||||
▼
|
||||
S3-b/c/d (各服务 CLI 适配)
|
||||
└─▶ S3-e (默认值注入) ──▶ Step 3 完成
|
||||
│
|
||||
▼
|
||||
S4 (端到端测试)
|
||||
```
|
||||
|
||||
### 7.3 新增/改动文件清单
|
||||
|
||||
```
|
||||
core/
|
||||
├── include/registry/
|
||||
│ └── InitParser.hh # 【新增】YAML/TXT 解析器头文件
|
||||
├── src/registry/
|
||||
│ └── InitParser.cpp # 【新增】YAML/TXT 解析器实现
|
||||
├── src/main/pmc_mtd.cpp # 【改动】调用 InitParser 替代 parse_self_init_list
|
||||
├── src/main.cpp # 【改动】--init 帮助文本更新
|
||||
├── Makefile # 【改动】追加 InitParser.o 编译目标 + -lyaml-cpp 链接
|
||||
├── _run.sh # 【改动】-i init.txt → -i .(目录模式自动查找)
|
||||
├── services.yaml # 【新增】声明式服务启动配置
|
||||
└── init.txt # 【保留】向后兼容回退
|
||||
|
||||
项目根目录/
|
||||
├── _setup.deb.sh # 【改动】追加 libyaml-cpp-dev
|
||||
└── _setup.rhel.sh # 【改动】追加 yaml-cpp-devel
|
||||
|
||||
各服务 _run.sh(11 个) # 【已完成】Step 2 统一 "$@" 透传 + 移除 sleep
|
||||
|
||||
C++ 服务源码(6 个) # 【待改】Step 3 新增/扩展 program_options
|
||||
Python 服务源码(4 个) # 【待改】Step 3 新增 argparse
|
||||
Go 服务源码(1 个) # 【待改】Step 3 扩展 flag
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 8. 验证清单
|
||||
|
||||
### 8.1 Step 1 验证
|
||||
|
||||
- [ ] `make` 编译成功,yaml-cpp 源文件正确编译链接
|
||||
- [ ] `services.yaml` 存在时,core 按 YAML 声明启动全部 11 个服务
|
||||
- [ ] `services.yaml` 不存在时,core 回退 `init.txt`,行为与改造前完全一致
|
||||
- [ ] `args` 字段的参数正确透传到 `_run.sh`
|
||||
- [ ] `env` 字段的环境变量在子进程启动前生效
|
||||
- [ ] `sleep_before` 字段正确控制启动间隔
|
||||
|
||||
### 8.2 Step 2 验证
|
||||
|
||||
- [x] 全部 11 个 `_run.sh` 末尾 `exec` 行带 `"$@"`
|
||||
- [x] waf 保留 `-c config.yaml` 作为脚本默认值,追加 `"$@"` 透传,services.yaml `args: []`
|
||||
- [x] aibox 的 `$@` 改为 `"$@"`(加引号,防空格路径问题)
|
||||
- [x] ca/gateway/captcha2/smtp 的 `sleep` 从 `_run.sh` 移除,由 `services.yaml` 的 `sleep_before` 管理
|
||||
- [ ] 手动执行 `_run.sh --help` 能透传到二进制(依赖 Step 3 CLI 适配)
|
||||
|
||||
### 8.3 Step 3 验证
|
||||
|
||||
- [ ] 全部服务支持 `--help` 输出用法
|
||||
- [ ] 全部服务支持 `--version` 输出版本
|
||||
- [ ] `--port` 能覆盖默认端口并生效
|
||||
- [ ] `--log-level` 能控制日志输出级别
|
||||
- [ ] `_run.sh` 默认值注入 + `services.yaml` 覆盖,两者叠加不冲突
|
||||
- [ ] ca/gateway 的环境变量桥接模式正确工作
|
||||
|
||||
---
|
||||
|
||||
## 9. 风险与缓解
|
||||
|
||||
| 风险 | 影响 | 缓解措施 |
|
||||
|------|------|----------|
|
||||
| yaml-cpp 编译增加构建时间 | 低(~10s) | 可接受;如需优化可预编译为静态库 |
|
||||
| pmc_mtd 改造影响核心启动路径 | 高 | 保留 `init.txt` 回退路径;先在测试环境验证 |
|
||||
| _run.sh 移除 sleep 后启动时序变化 | 中 | `services.yaml` 的 `sleep_before` 提供等价等待 |
|
||||
| 各服务 CLI 参数适配工作量大 | 中 | 分批适配;_run.sh 默认值注入可过渡 |
|
||||
| Go flag 与 C++ program_options 语法差异 | 低 | 两者都支持 `-H value` 和 `--host value` 格式 |
|
||||
| Python argparse 与 C++ program_options 行为差异 | 低 | 测试覆盖;argparse 默认不覆盖已设默认值 |
|
||||
|
||||
---
|
||||
|
||||
## 10. 附录
|
||||
|
||||
### A. 各服务端口与配置对照
|
||||
|
||||
| 服务 | 端口 | 协议 | 配置来源 | CLI 框架现状 |
|
||||
|------|------|------|----------|-------------|
|
||||
| db2 | 8005 | pymsg | 硬编码 | 无 |
|
||||
| ai_review | 8007 | http | 硬编码 | 无 |
|
||||
| dnfs2 | 8006 | http | 硬编码 + confs | 无 |
|
||||
| smtp | 8008 | pymsg | 硬编码 | 无 |
|
||||
| fsys | 8090 | http | 硬编码 | 无 |
|
||||
| waf | 443 | https | YAML 配置文件 | Go flag ✅ |
|
||||
| ca | 8443 | https | confs + 环境变量 | program_options(仅 --help/--version) |
|
||||
| gateway | 8445 | http | confs + 环境变量 | program_options(仅 --help/--version) |
|
||||
| captcha2 | 8009 | pymsg | 硬编码 | 无 |
|
||||
| aibox | 8765 | websocket | _setenv.sh + 硬编码 | argparse(已 import 未使用) |
|
||||
| redirect80 | 80 | http | 硬编码 | 无 |
|
||||
|
||||
### B. yaml-cpp 集成验证
|
||||
|
||||
已采用系统包方式(非 vendored),在 `_setup.deb.sh` / `_setup.rhel.sh` 中声明依赖:
|
||||
|
||||
```bash
|
||||
# Debian/Ubuntu
|
||||
sudo apt install libyaml-cpp-dev
|
||||
|
||||
# RHEL/CentOS
|
||||
sudo dnf install yaml-cpp-devel
|
||||
```
|
||||
|
||||
Makefile 中 `-lyaml-cpp` 链接,`InitParser.cpp` 中 `#include <yaml-cpp/yaml.h>` 引用。
|
||||
|
||||
### C. 环境变量约定(现有 + 新增)
|
||||
|
||||
| 环境变量 | 使用方 | 当前/新增 | 说明 |
|
||||
|----------|--------|-----------|------|
|
||||
| `CONFS_HOST` | ca, gateway, dnfs2 | 现有 | confs 服务地址,默认 `127.0.0.1` |
|
||||
| `CONFS_PORT` | ca, gateway, dnfs2 | 现有 | confs 服务端口,默认 `8010` |
|
||||
| `CORE_HOST` | 未来 RegistrationClient | 新增 | core 注册中心地址,默认 `127.0.0.1` |
|
||||
| `CORE_PORT` | 未来 RegistrationClient | 新增 | core 注册中心端口,默认 `5200` |
|
||||
| `LOG_LEVEL` | 全部服务(可选) | 新增 | 日志级别,可被 `--log-level` 覆盖 |
|
||||
|
||||
### D. 与注册中心计划的衔接点
|
||||
|
||||
| 本方案产出 | 注册中心计划消费点 |
|
||||
|-----------|-------------------|
|
||||
| `services.yaml` 格式 | P7 直接复用,追加 `health_endpoint`/`protocol`/`restart_max` 等字段 |
|
||||
| `InitParser` 解析器 | P2 调用 `parseYaml()` 后自动填充 `ServiceRegistry` |
|
||||
| `ServiceConfig.depends_on` | P7 DAG 拓扑排序的输入 |
|
||||
| `ServiceConfig.args` | 各服务通过 CLI 参数接收注册中心分配的端口/实例名 |
|
||||
| CLI `--instance` 参数 | P5 RegistrationClient SDK 用此值作为注册服务名 |
|
||||
| CLI `--port` 参数 | P0 ServiceEntry 的 port 字段来源 |
|
||||
Loaded 100 of 140 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user