Compare commits
14
Commits
3e860c16e0
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2832566e87 | ||
|
|
7d5d7b9e36 | ||
|
|
4497d80a94 | ||
|
|
04f1375b45 | ||
|
|
d9b2359d25 | ||
|
|
cc72248e26 | ||
|
|
c4154d4db9 | ||
|
|
0ed03ac095 | ||
|
|
2ea3fa14d8 | ||
|
|
dbfb925f65 | ||
|
|
67acae708c | ||
|
|
7e44025671 | ||
|
|
2dc79f0637 | ||
|
|
5894e64287 |
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/
|
||||
+5
-5
@@ -109,7 +109,7 @@ aibox/
|
||||
│ ├── proxy_tools.py # B 批:SetProxy/GetProxy
|
||||
│ ├── misc_tools.py # B 批:AskUserTool
|
||||
│ └── process_tools.py # 后台进程:BgRun/BgWrite/BgList/BgOutput/BgKill(bg_run/bg_write/bg_list/bg_output/bg_kill)
|
||||
├── tests/ # 单元测试套件(529 项全过)
|
||||
├── tests/ # 单元测试套件(598 项全过)
|
||||
│ ├── conftest.py # sys.path 配置 + websockets 桩
|
||||
│ ├── test_basic_tools.py # A 批工具 43 项
|
||||
│ ├── test_b_tools.py # B 批工具 58 项
|
||||
@@ -164,7 +164,7 @@ common ← subsystems ← tools ← aibox.py
|
||||
```bash
|
||||
# 从仓库根目录运行
|
||||
py -m pytest aibox/tests -q
|
||||
# 当前:529 项全过(Linux;test_execute_timeout / test_flat_api_still_reads_key_file 为 Windows 向用例,Linux 预存失败)
|
||||
# 当前:598 项全过(Linux;test_execute_timeout 为 Windows 向用例,本机无 ping 预存失败)
|
||||
# Windows 开发机:532 passed / 4 failed(平台环境项:symlink 权限、python3 缺失、POSIX 可执行位语义 ×2)/ 1 skipped
|
||||
```
|
||||
|
||||
@@ -247,7 +247,7 @@ py -m pytest aibox/tests -q
|
||||
| `bg_output` | 拉取输出:`tail=N` 尾部 N 行 / `since=游标` 增量续读(游标为上次返回的 `next_cursor`);单次默认 64KB、上限 1MB;环形缓冲(默认 256KB/进程)淘汰的早期行如实提示 |
|
||||
| `bg_kill` | 终止后台进程(整棵进程组):三级优雅终止——关 stdin(行协议进程见 EOF 常自行退出)→ SIGTERM 宽限 → SIGKILL 兜底;`force=true` 直接 SIGKILL;对已退出进程幂等 |
|
||||
|
||||
> **后台进程语义(v3.8)**:注册表为全局单例(跨会话/跨子代理可见,任何 worker 可管理其他 worker 启动的进程);用户中断(Ctrl+C/取消事件)**不**影响后台进程;aibox 进程退出经 atexit 统一收尾(非 detach 全杀,SIGKILL/SIGTERM 硬杀不覆盖);单行输出截断 16KB;运行中上限 16 / 总条目上限 32 / 单次 stdin 写入上限 64KB(`background_processes` 配置节可调);POSIX 按进程组终止(`start_new_session`,孙进程不悬挂),Windows 回退单进程终止。
|
||||
> **后台进程语义(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
|
||||
|
||||
@@ -934,11 +934,11 @@ 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` 豁免但脱离句柄管理);注册表跨会话/跨子代理共享(任何 worker 可管理其他 worker 启动的进程)
|
||||
- **生命周期语义(三决策)**:用户中断(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 回退单进程终止;TTL `max_lifetime` 惰性清理(默认 0 不启用);凭据类交互(sudo -S 喂密码)会留对话记录,敏感凭据走环境变量注入(工具描述已指引)
|
||||
- **安全边界**:运行中上限 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(平台环境项,与代码无关)
|
||||
|
||||
|
||||
+95
-29
@@ -121,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):
|
||||
@@ -318,14 +323,24 @@ 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)
|
||||
@@ -350,8 +365,22 @@ class TaskScheduler:
|
||||
|
||||
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:
|
||||
if timed_out.is_set():
|
||||
self.logger.error(f"任务执行超时({task_timeout}秒),AI 已取消")
|
||||
@@ -987,6 +1016,10 @@ 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",
|
||||
@@ -1025,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
|
||||
@@ -1182,18 +1219,22 @@ class DeepSeekChat:
|
||||
SimpleWebCrawlerTool(self.logger),
|
||||
CalculatorTool(self.logger),
|
||||
GetCurrentTimeTool(self.logger),
|
||||
# 后台进程管理(共享全局单例:跨子代理可管理彼此启动的进程)
|
||||
BgRunTool(self.logger),
|
||||
BgWriteTool(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),
|
||||
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),
|
||||
]
|
||||
@@ -1204,11 +1245,12 @@ 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:
|
||||
@@ -1224,7 +1266,9 @@ class DeepSeekChat:
|
||||
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))
|
||||
|
||||
# 子代理工具:主代理可派生独立上下文/角色的轻量子代理执行任务
|
||||
@@ -1376,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
|
||||
@@ -1414,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 中写明执行者身份,"
|
||||
@@ -1423,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 判定、
|
||||
@@ -1430,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
|
||||
@@ -1448,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")
|
||||
@@ -3128,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))
|
||||
|
||||
|
||||
+18
-1
@@ -10,6 +10,7 @@ mkstemp 随机文件名避免多进程并发写同路径互踩(conversation_hi
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
|
||||
|
||||
def atomic_write_text(path: str, content: str) -> None:
|
||||
@@ -23,10 +24,26 @@ def atomic_write_text(path: str, content: str) -> None:
|
||||
f.write(content)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
os.replace(tmp_path, path)
|
||||
# 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
|
||||
|
||||
+85
-20
@@ -29,13 +29,28 @@ 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 令牌。"""
|
||||
"""日志脱敏:替换错误消息中的 Bearer 令牌与裸 API Key。"""
|
||||
if not msg:
|
||||
return msg
|
||||
return _AUTH_HEADER_RE.sub("Bearer ***", msg)
|
||||
msg = _AUTH_HEADER_RE.sub("Bearer ***", msg)
|
||||
msg = _API_KEY_RE.sub("sk-***", msg)
|
||||
return msg
|
||||
|
||||
|
||||
class LLMClient(ABC):
|
||||
@@ -166,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
|
||||
@@ -219,14 +243,22 @@ 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:
|
||||
@@ -236,6 +268,7 @@ class DeepSeekClient(LLMClient):
|
||||
continue
|
||||
data = line_text[6:]
|
||||
if data == "[DONE]":
|
||||
saw_done = True
|
||||
yield {"type": "done"}
|
||||
break
|
||||
try:
|
||||
@@ -251,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")
|
||||
@@ -368,8 +409,17 @@ class FailoverClient:
|
||||
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
|
||||
@@ -377,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:
|
||||
@@ -434,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
|
||||
@@ -451,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):
|
||||
@@ -474,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:
|
||||
|
||||
+43
-5
@@ -282,16 +282,49 @@ def _choose(prompt: str, options: List[str], multi: bool = False) -> List[str]:
|
||||
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"环境变量 {p['api_key_env']}" if p["api_key_env"] else "免密"
|
||||
fb = f",失败转移 → {', '.join(p['fallback'])}" if p["fallback"] else ""
|
||||
print(f" {p['name']}{mark} [{p['display_name']}]")
|
||||
print(f" 地址: {p['api_url']} 模型: {p['model']}")
|
||||
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}")
|
||||
|
||||
|
||||
@@ -319,7 +352,12 @@ def _add_flow(config: ConfigManager):
|
||||
key_mode = _ask("选择 [1/2]", "1")
|
||||
api_key_env = ""
|
||||
if key_mode == "1":
|
||||
api_key_env = _ask("环境变量名", "DEEPSEEK_API_KEY")
|
||||
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
|
||||
|
||||
@@ -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)
|
||||
|
||||
+16
-7
@@ -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,7 +148,7 @@ 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)删除索引条目需从内容表回读旧值,内容行
|
||||
@@ -166,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()
|
||||
|
||||
@@ -240,7 +249,7 @@ class MemoryDatabase:
|
||||
的 None→默认上限语义对齐)。"""
|
||||
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()
|
||||
cursor.execute('SELECT * FROM memories ORDER BY end_time DESC LIMIT ?', (limit,))
|
||||
@@ -256,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,))
|
||||
@@ -274,7 +283,7 @@ 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 再删内容表,防索引条目残留
|
||||
|
||||
+149
-54
@@ -22,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):
|
||||
"""命令执行器抽象基类 - 负责执行系统命令"""
|
||||
@@ -71,30 +101,82 @@ 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 _terminate_tree(self, proc, force: bool = False):
|
||||
def _register_proc(self, process):
|
||||
"""T5:注册本次执行的进程槽,返回 (token, pgid)"""
|
||||
with self._process_lock:
|
||||
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,届时无法再从进程反查组。
|
||||
Windows 无进程组概念,退回单进程 terminate/kill。"""
|
||||
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
|
||||
with self._process_lock:
|
||||
pgid = self._current_pgid
|
||||
if pgid is None:
|
||||
pgid = os.getpgid(proc.pid)
|
||||
os.killpg(pgid, sig)
|
||||
elif force:
|
||||
proc.kill()
|
||||
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:
|
||||
proc.terminate()
|
||||
# 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:
|
||||
proc.kill()
|
||||
@@ -102,18 +184,21 @@ class SubprocessCommandExecutor(CommandExecutor):
|
||||
pass
|
||||
|
||||
def cancel(self):
|
||||
"""取消当前正在运行的命令(kill子进程组)。线程安全。"""
|
||||
"""取消本执行器当前注册的全部运行中命令(kill 子进程组)。线程安全。
|
||||
T5:多槽令牌注册表——并发执行时定向终止各自槽位,互不误杀;
|
||||
无参调用语义为"终止本执行器全部在途命令"(中断/关闭路径)。"""
|
||||
with self._process_lock:
|
||||
proc = self._current_process
|
||||
if proc is not None:
|
||||
slots = list(self._proc_slots.values())
|
||||
for proc, pgid in slots:
|
||||
try:
|
||||
if proc.poll() is None:
|
||||
self._terminate_tree(proc, force=True)
|
||||
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) -> Tuple[str, bool, bool]:
|
||||
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 线程执行,主循环带超时轮询队列。
|
||||
@@ -147,14 +232,14 @@ class SubprocessCommandExecutor(CommandExecutor):
|
||||
# J16(审计 2026-08-19):timeout<=0 语义为"不设超时"(有意为之,
|
||||
# 调用方显式传 0 禁用);registry 层 3600s 流式兜底仍生效
|
||||
if timeout > 0 and time.time() - start_time > timeout:
|
||||
self._terminate_tree(process)
|
||||
self._terminate_tree(process, pgid=pgid)
|
||||
timeout_occurred = True
|
||||
break
|
||||
|
||||
# 检查输出大小
|
||||
if total_bytes >= self.MAX_OUTPUT_BYTES:
|
||||
truncated = True
|
||||
self._terminate_tree(process)
|
||||
self._terminate_tree(process, pgid=pgid)
|
||||
break
|
||||
|
||||
try:
|
||||
@@ -175,7 +260,7 @@ class SubprocessCommandExecutor(CommandExecutor):
|
||||
# 与管道 fd 至孙进程退出,写满 64KB 缓冲后孙进程自身
|
||||
# 也卡死。此处强杀进程组收尾(nohup 重定向输出的
|
||||
# 真正脱离场景不持管道,不会走到这里)。
|
||||
self._terminate_tree(process, force=True)
|
||||
self._terminate_tree(process, force=True, pgid=pgid)
|
||||
break
|
||||
continue
|
||||
|
||||
@@ -188,7 +273,7 @@ class SubprocessCommandExecutor(CommandExecutor):
|
||||
output_chunks.append(line.encode('utf-8', errors='replace')[:allowed].decode('utf-8', errors='ignore'))
|
||||
total_bytes = self.MAX_OUTPUT_BYTES
|
||||
truncated = True
|
||||
self._terminate_tree(process)
|
||||
self._terminate_tree(process, pgid=pgid)
|
||||
break
|
||||
|
||||
output_chunks.append(line)
|
||||
@@ -224,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,
|
||||
@@ -246,31 +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),
|
||||
# L9:POSIX 下 shell 成为会话首进程,超时/取消可按进程组
|
||||
# 一次终止整棵进程树(含 & 后台孙进程),消除管道悬挂
|
||||
start_new_session=(os.name == "posix")
|
||||
)
|
||||
# 保存进程引用以便cancel()可以kill它
|
||||
with self._process_lock:
|
||||
self._current_process = process
|
||||
if os.name == "posix":
|
||||
try:
|
||||
self._current_pgid = os.getpgid(process.pid)
|
||||
except Exception:
|
||||
self._current_pgid = None
|
||||
# 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:
|
||||
self._terminate_tree(process, force=timeout_occurred or truncated)
|
||||
self._terminate_tree(process, force=timeout_occurred or truncated, pgid=pgid)
|
||||
try:
|
||||
process.wait(timeout=2)
|
||||
except:
|
||||
self._terminate_tree(process, force=True)
|
||||
self._terminate_tree(process, force=True, pgid=pgid)
|
||||
|
||||
return_code = process.returncode if process.returncode is not None else -1
|
||||
|
||||
@@ -298,11 +393,10 @@ class SubprocessCommandExecutor(CommandExecutor):
|
||||
return f"❌ 执行错误: {str(e)}"
|
||||
finally:
|
||||
if process and process.poll() is None:
|
||||
self._terminate_tree(process, force=True)
|
||||
# 清除当前进程引用
|
||||
with self._process_lock:
|
||||
self._current_process = None
|
||||
self._current_pgid = 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):
|
||||
@@ -323,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,
|
||||
@@ -339,6 +438,8 @@ 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),
|
||||
# L9:POSIX 下 shell 成为会话首进程,超时/取消可按进程组
|
||||
@@ -346,19 +447,13 @@ class SubprocessCommandExecutor(CommandExecutor):
|
||||
start_new_session=(os.name == "posix")
|
||||
)
|
||||
self.process = process
|
||||
# 保存进程引用以便cancel()可以kill它
|
||||
with self._process_lock:
|
||||
self._current_process = process
|
||||
if os.name == "posix":
|
||||
try:
|
||||
self._current_pgid = os.getpgid(process.pid)
|
||||
except Exception:
|
||||
self._current_pgid = None
|
||||
# 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"))
|
||||
@@ -367,7 +462,7 @@ class SubprocessCommandExecutor(CommandExecutor):
|
||||
|
||||
# 如果超时或截断,确保进程被终止
|
||||
if (timeout_occurred or truncated) and process.poll() is None:
|
||||
self._terminate_tree(process, force=True)
|
||||
self._terminate_tree(process, force=True, pgid=pgid)
|
||||
try:
|
||||
process.wait(timeout=2)
|
||||
except:
|
||||
@@ -404,12 +499,12 @@ class SubprocessCommandExecutor(CommandExecutor):
|
||||
output_queue.put(("complete", f"❌ 执行错误: {str(e)}"))
|
||||
finally:
|
||||
if process and process.poll() is None:
|
||||
self._terminate_tree(process)
|
||||
self._terminate_tree(process, pgid=pgid)
|
||||
try:
|
||||
process.wait(timeout=2)
|
||||
except:
|
||||
self._terminate_tree(process, force=True)
|
||||
self._terminate_tree(process, force=True, pgid=pgid)
|
||||
self.process = None
|
||||
with self._process_lock:
|
||||
self._current_process = None
|
||||
self._current_pgid = None
|
||||
# T5:仅注销自己的槽位,不再无差别清空共享引用
|
||||
if token is not None:
|
||||
self._unregister_proc(token)
|
||||
@@ -35,6 +35,7 @@ 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
|
||||
@@ -76,7 +77,7 @@ class _ProcEntry:
|
||||
self.max_buffer = max(1024, int(buffer_bytes))
|
||||
self.next_seq = 1
|
||||
self.evicted = 0
|
||||
self.lock = threading.Lock()
|
||||
self.lock = threading.RLock() # T8:可重入,write_input 持锁期间可调用 _record_line
|
||||
self.reader_done = threading.Event()
|
||||
self.reader = None
|
||||
# 交互写入:单写锁保序(bg_write 串行化,防交错输入)
|
||||
@@ -162,7 +163,11 @@ class ProcessManager:
|
||||
|
||||
exec_env = os.environ.copy()
|
||||
if env:
|
||||
exec_env.update(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(
|
||||
@@ -178,6 +183,8 @@ class ProcessManager:
|
||||
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 按进程组
|
||||
@@ -320,7 +327,7 @@ class ProcessManager:
|
||||
|
||||
def write(self, handle, text: str, append_newline: bool = True,
|
||||
wait_lines: int = 0, wait_timeout: float = 5.0,
|
||||
write_timeout: float = 5.0) -> Dict:
|
||||
write_timeout: float = 5.0, owner: str = None) -> Dict:
|
||||
"""向后台进程 stdin 写入一行(交互式问答)。
|
||||
|
||||
返回 {ok, handle, watermark, sent_bytes[, wait]}:
|
||||
@@ -328,10 +335,16 @@ class ProcessManager:
|
||||
watermark) 得到的必然且仅是本条输入引发的响应(写读对齐)
|
||||
- wait_lines>0 时同步等待 N 行响应(wait 子 dict 含 output/
|
||||
timed_out),expect 式一步问答
|
||||
失败返回 {ok: False, error}(进程已退出/写入超时/超长等)。"""
|
||||
失败返回 {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"
|
||||
@@ -345,12 +358,13 @@ class ProcessManager:
|
||||
if entry.status != "running":
|
||||
return {"ok": False, "error": f"进程已结束({entry.status}),无法写入"}
|
||||
|
||||
# watermark:回显行将占用的序号(写入前先记账,响应必然在其后)
|
||||
# watermark:取号与回显行记账必须在同一临界区,保证 watermark 等于
|
||||
# 回显行真实序号,bg_output(since=watermark) 正确对齐本条响应(T8)
|
||||
echo = "[stdin] " + text.replace("\n", "⏎")
|
||||
with entry.lock:
|
||||
watermark = entry.next_seq
|
||||
# 回显行进环形缓冲 + log 镜像(多行输入折叠为单行展示)
|
||||
echo = "[stdin] " + text.replace("\n", "⏎")
|
||||
self._record_line(entry, echo)
|
||||
# 回显行进环形缓冲 + log 镜像(多行输入折叠为单行展示)
|
||||
self._record_line(entry, echo)
|
||||
|
||||
# 写入放 watchdog 线程:进程不读 stdin 且管道缓冲(~64KB)写满时
|
||||
# write() 会阻塞,不能拖死工具线程;超时后 daemon 线程由系统兜底
|
||||
@@ -464,12 +478,20 @@ class ProcessManager:
|
||||
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) -> Dict:
|
||||
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)
|
||||
@@ -559,6 +581,20 @@ class ProcessManager:
|
||||
# 内部实现
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
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)
|
||||
@@ -599,7 +635,7 @@ class ProcessManager:
|
||||
self.logger.info(f"后台进程 #{entry.handle} TTL 到期已终止")
|
||||
|
||||
def _terminate_tree(self, entry: _ProcEntry, force: bool = False):
|
||||
"""按进程组终止整棵进程树(对齐 executors L9 语义与 Windows 回退)"""
|
||||
"""按进程组终止整棵进程树(对齐 executors L9/T11 语义与 Windows 回退)"""
|
||||
proc = entry.popen
|
||||
try:
|
||||
if os.name == "posix":
|
||||
@@ -608,10 +644,21 @@ class ProcessManager:
|
||||
if pgid is None:
|
||||
pgid = os.getpgid(proc.pid)
|
||||
os.killpg(pgid, sig)
|
||||
elif force:
|
||||
proc.kill()
|
||||
else:
|
||||
proc.terminate()
|
||||
# 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()
|
||||
@@ -685,6 +732,7 @@ class ProcessManager:
|
||||
"buffered_lines": len(entry.lines),
|
||||
"evicted_lines": entry.evicted,
|
||||
"command": entry.command,
|
||||
"owner": entry.owner, # T22:属主标识(观察面可见,写/杀需匹配)
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -123,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 = [
|
||||
@@ -153,7 +159,7 @@ class SkillManager:
|
||||
is_markdown = self._is_markdown_file(skill_file)
|
||||
|
||||
if is_markdown:
|
||||
if skill_name in self.skills:
|
||||
if skill_name in new_skills:
|
||||
self.logger.warning(
|
||||
f"Markdown 技能 '{skill_file}' 与同名可执行技能冲突,已跳过"
|
||||
f"(可执行技能优先)")
|
||||
@@ -161,7 +167,7 @@ class SkillManager:
|
||||
# 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,
|
||||
@@ -182,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,
|
||||
@@ -192,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:
|
||||
"""重新扫描技能目录,返回加载的技能数量。"""
|
||||
|
||||
@@ -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 # 可打印部分保留
|
||||
+180
-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):
|
||||
@@ -106,9 +142,54 @@ class ReadFileTool(FileToolBase):
|
||||
return "❌ 参数错误: start_line/end_line 必须为整数"
|
||||
start_line = max(1, start_line)
|
||||
end_line = max(0, min(end_line, 10_000_000))
|
||||
quoted = shlex.quote(abs_path)
|
||||
|
||||
# 构造等价的命令行读取命令
|
||||
# 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:
|
||||
@@ -180,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 ""
|
||||
|
||||
@@ -254,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 提示词,指导智能体用命令行工具列目录
|
||||
@@ -324,7 +457,9 @@ 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"):
|
||||
@@ -336,10 +471,44 @@ class SimpleWebCrawlerTool(Tool):
|
||||
if not url or not str(url).strip().lower().startswith(("http://", "https://")):
|
||||
return "❌ 仅支持 http/https 协议地址(已拦截 file:// 等非网络协议)"
|
||||
|
||||
quoted_url = shlex.quote(url)
|
||||
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":
|
||||
@@ -358,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 完成网页访问
|
||||
|
||||
@@ -78,9 +78,11 @@ def _coerce_bool(value, default: bool) -> bool:
|
||||
class BgRunTool(Tool):
|
||||
"""后台启动长驻进程"""
|
||||
|
||||
def __init__(self, manager: ProcessManager = None, logger=None):
|
||||
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"
|
||||
@@ -133,8 +135,14 @@ class BgRunTool(Tool):
|
||||
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)
|
||||
env=env, detach=detach, log_file=log_file,
|
||||
owner=self.owner)
|
||||
if not result.get("ok"):
|
||||
return self.format_result(False, "后台进程启动失败",
|
||||
{"错误": result.get("error", "未知错误")})
|
||||
@@ -153,9 +161,11 @@ class BgRunTool(Tool):
|
||||
class BgWriteTool(Tool):
|
||||
"""向后台进程 stdin 写入(交互式问答)"""
|
||||
|
||||
def __init__(self, manager: ProcessManager = None, logger=None):
|
||||
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"
|
||||
@@ -215,6 +225,7 @@ class BgWriteTool(Tool):
|
||||
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} 失败",
|
||||
@@ -352,9 +363,11 @@ class BgOutputTool(Tool):
|
||||
class BgKillTool(Tool):
|
||||
"""终止后台进程"""
|
||||
|
||||
def __init__(self, manager: ProcessManager = None, logger=None):
|
||||
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"
|
||||
@@ -393,7 +406,10 @@ class BgKillTool(Tool):
|
||||
if handle is None:
|
||||
return "❌ 错误:handle 必须是整数句柄"
|
||||
|
||||
result = self.manager.kill(handle, force=force, grace=grace)
|
||||
# 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", "未知错误")})
|
||||
|
||||
+30
-12
@@ -68,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:
|
||||
"""执行工具"""
|
||||
@@ -85,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()
|
||||
@@ -134,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
|
||||
@@ -146,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
|
||||
|
||||
@@ -174,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
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
+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;
|
||||
@@ -1,9 +1,9 @@
|
||||
# aibox 模块安全与质量审计报告
|
||||
|
||||
- **日期**:2026-08-17(第一轮)/ 2026-08-18(第二轮复审)/ 2026-08-18(第三轮:遗留复核 + v3.4 专审)/ 2026-08-18(第四轮:3 路并行审读 + Batch H 修复)/ 2026-08-19(第五轮:4 路并行审读,聚焦历届未覆盖区域)/ 2026-08-20(第六轮:交叉视角——修复回归验证 / 并发与生命周期 / 降级路径 / 修复间相互作用)/ 2026-08-20(第七轮:skills 全量脚本 / 前端深维度 / 客户端三文件 / WS+CLI 命令面全量校验)
|
||||
- **审计基线**:第一轮 HEAD `ce142fe`(Batch A–D 后)→ 第二轮 HEAD `177a9ac`(Batch E 修复后)→ 第三轮 HEAD `aa15841`(L4a/b/c 落地后)→ 第四轮 HEAD `d1f780e`(Batch G + SOCKS5 验证后)→ 第五轮 HEAD `d513c5f`(Batch I + model_config.py 后,含工作区未提交的 max_turns 钳制)→ 第六轮 HEAD `7dc6796`(Batch K 全落地 + 五轮收束文档同步后)→ 第七轮 HEAD `e43aeea`(Batch L 落地后)
|
||||
- **审计范围**:`aibox.py`(~3458 行)+ `llm_client.py` + `model_config.py` + `tools/`(7 文件)+ `subsystems/`(6 文件)+ `common/`(5 文件)+ 前端 HTML(frontend.html / entry.html / entry1.html)+ `skills/` 脚本 + `client.py` / `voice_client.py`(前端与技能脚本为第五轮首次覆盖;第七轮完成 skills 全量逐文件与命令面逐条覆盖)
|
||||
- **方法**:第一轮逐文件代码审读 + PoC 动态验证;第二轮 Batch E 修复复审 + 全量重审 + 6 处关键发现实证抽查;第三轮遗留项复核 + v3.4 专审;第四轮 3 路并行审读 + 主审逐条复核 + Linux PoC 实证(6/6);第五轮 4 路并行审读(聚焦历届未覆盖区域)+ 主审逐条复核 + Linux PoC 实证(2/2);第六轮 4 路并行审读(Batch K 修复回归验证与修复自身缺陷面 / 并发与资源生命周期交叉 / 数据边界与故障降级 / 工具链与修复间相互作用)+ 主审逐条代码复核(Windows 开发机,无 Linux PoC,动态行为以代码级核验 + Windows 测试基线代替);第七轮 4 路并行审读(skills 全量安全面 / 前端深维度 / model_config+voice_client+client 深审 / WS+CLI 命令面全量逐条校验)+ 主审对全部中危项逐条代码实证(Windows 开发机)
|
||||
- **日期**:2026-08-17(第一轮)/ 2026-08-18(第二轮复审)/ 2026-08-18(第三轮:遗留复核 + v3.4 专审)/ 2026-08-18(第四轮:3 路并行审读 + Batch H 修复)/ 2026-08-19(第五轮:4 路并行审读,聚焦历届未覆盖区域)/ 2026-08-20(第六轮:交叉视角——修复回归验证 / 并发与生命周期 / 降级路径 / 修复间相互作用)/ 2026-08-20(第七轮:skills 全量脚本 / 前端深维度 / 客户端三文件 / WS+CLI 命令面全量校验)/ 2026-08-21(第八轮:后台进程池 v3.8 新代码首审 + 全模块 3 路并行深审)
|
||||
- **审计基线**:第一轮 HEAD `ce142fe`(Batch A–D 后)→ 第二轮 HEAD `177a9ac`(Batch E 修复后)→ 第三轮 HEAD `aa15841`(L4a/b/c 落地后)→ 第四轮 HEAD `d1f780e`(Batch G + SOCKS5 验证后)→ 第五轮 HEAD `d513c5f`(Batch I + model_config.py 后,含工作区未提交的 max_turns 钳制)→ 第六轮 HEAD `7dc6796`(Batch K 全落地 + 五轮收束文档同步后)→ 第七轮 HEAD `e43aeea`(Batch L 落地后)→ 第八轮 HEAD `3e860c1`(Batch M 全 5 波落地 + 后台进程池 v3.8,commit 5cb3594)
|
||||
- **审计范围**:`aibox.py`(~3408 行)+ `llm_client.py` + `model_config.py` + `tools/`(10 文件,含 v3.8 新增 `process_tools.py`)+ `subsystems/`(7 文件,含 v3.8 新增 `process_manager.py`)+ `common/`(5 文件)+ 前端 HTML(frontend.html / entry.html / entry1.html)+ `skills/` 脚本 + `client.py` / `voice_client.py`(前端与技能脚本为第五轮首次覆盖;第七轮完成 skills 全量逐文件与命令面逐条覆盖;第八轮首次覆盖后台进程池子系统)
|
||||
- **方法**:第一轮逐文件代码审读 + PoC 动态验证;第二轮 Batch E 修复复审 + 全量重审 + 6 处关键发现实证抽查;第三轮遗留项复核 + v3.4 专审;第四轮 3 路并行审读 + 主审逐条复核 + Linux PoC 实证(6/6);第五轮 4 路并行审读(聚焦历届未覆盖区域)+ 主审逐条复核 + Linux PoC 实证(2/2);第六轮 4 路并行审读(Batch K 修复回归验证与修复自身缺陷面 / 并发与资源生命周期交叉 / 数据边界与故障降级 / 工具链与修复间相互作用)+ 主审逐条代码复核(Windows 开发机,无 Linux PoC,动态行为以代码级核验 + Windows 测试基线代替);第七轮 4 路并行审读(skills 全量安全面 / 前端深维度 / model_config+voice_client+client 深审 / WS+CLI 命令面全量逐条校验)+ 主审对全部中危项逐条代码实证(Windows 开发机);第八轮 3 路并行深审(aibox.py 编排层 / llm_client+model_config / subsystems+tools+common)+ 主审对 v3.8 新代码(process_manager.py/process_tools.py)逐行复核 + Windows 测试基线确认
|
||||
|
||||
### 轮次总览
|
||||
|
||||
@@ -15,7 +15,8 @@
|
||||
| 第四轮 | 2026-08-18 | `d1f780e` | Batch H 全项(含 N5 取消传播回归,高危) | Batch H + Batch I 遗留收尾 | 424 passed |
|
||||
| 第五轮 | 2026-08-19 | `d513c5f` | J1~J16(高危 2 / 中危 5 / 低危与信息 9) | Batch K(波次 1~5 全落地) | 490 passed |
|
||||
| 第六轮 | 2026-08-20 | `7dc6796` | R1~R3(均低危;Batch K 回归验证零回归) | Batch L(R1~R3 全落地) | 490 passed(Linux)/ 493 passed(Windows) |
|
||||
| 第七轮 | 2026-08-20 | `e43aeea` | 46 项(高危 0 / 中危 11 / 低危 26 / 信息 9;Batch L 回归验证零回归) | Batch M 波次 1 已落地(S1/S4/S5/S7/S8),波次 2~5 待落地 | 493 passed / 4 failed(平台环境项)/ 1 skipped(Windows) |
|
||||
| 第七轮 | 2026-08-20 | `e43aeea` | 46 项(高危 0 / 中危 11 / 低危 26 / 信息 9;Batch L 回归验证零回归) | Batch M 波次 1~5 全落地 | 493 passed / 4 failed(平台环境项)/ 1 skipped(Windows) |
|
||||
| 第八轮 | 2026-08-21 | `3e860c1` | T1~T11 高危 11 / T12~T28 中危 17 / 低危与信息 20+(Batch M 全量回归验证;v3.8 后台进程池首审) | 待落地(Batch N 建议见 §23.7) | 532 passed / 4 failed(平台环境项)/ 1 skipped(Windows) |
|
||||
|
||||
---
|
||||
|
||||
@@ -35,7 +36,9 @@
|
||||
|
||||
**第七轮**(基线 e43aeea):四路新维度——skills 全量 29 脚本逐文件(继 J12 点修后首次)、前端三页面深维度、model_config/voice_client/client 深审、WS+CLI 命令面逐条校验;主审对全部中危项逐条代码实证。Batch L 零回归;新发现 **46 项**(高危 0 / 中危 11 / 低危 26 / 信息 9),集中于 skills 凭据进命令行(S1/S4/S5/S7,违反项目自定 J12 条例)、skills 路径边界(S2/S9/S11/S12,含 vision.py 任意文件经云 API 渗出的提示注入通道)、前端功能缺陷(F1 清空按钮必抛 TypeError / F2 手动断开被重连复活)、voice_client input 线程泄漏(M1)与命令面输入规整(W 系列)。Batch M 修复建议(5 波次)已给出,待落地。见"十七、第七轮审计"。
|
||||
|
||||
**最终状态**:七轮累计发现 118 项(高危 7 / 中危 37 / 低危与信息 74);前六轮 72 项**全部修复落地**(Batch E/F/G/H/I/K/L),第七轮 **Batch M 波次 1 已落地**(S1/S4/S5/S7/S8 凭据收口,见"十八"),波次 2~5(41 项,中危 6 / 低危与信息 35)待落地;维持原判仅限设计决策项,见"十五、最终状态与遗留清单"与"十六、第六轮审计"。
|
||||
**第八轮**(基线 3e860c1):Batch M 全 5 波落地后的首次全量深审,同时首审 v3.8 后台进程池(`process_manager.py` 731 行 + `process_tools.py` 408 行,commit 5cb3594)。方法为 3 路并行深审(aibox.py 编排层 / llm_client+model_config / subsystems+tools+common)+ 主审对 v3.8 新代码逐行复核。Batch M 零回归;新发现 **T1~T28 中高危 28 项 + 低危与信息 20+**,其中高危 11 项集中于三个系统性主题:**平台矩阵未声明**(POSIX 专用逻辑在 Windows 下失效或成注入面,T3/T11/T16)、**"单会话串行"隐含假设在多连接 WS 下不成立**(registry 输出队列串扰 / executor 单槽误杀 / watermark 竞态,T4/T5/T8)、**防线位置错位**(脱敏建在消费侧 / 冷却键与去重键不一致 / 生成器延迟副作用,T6/T7/T13)。本轮发现的 T2(set_proxy 全局网络面)、T1(WS 无鉴权)、T10(调度器无硬超时)为设计层问题,需产品决策。见"二十三、第八轮审计"。
|
||||
|
||||
**最终状态**:八轮累计发现 166+ 项(高危 18 / 中危 54 / 低危与信息 94+);前七轮 118 项**全部修复落地**(Batch E/F/G/H/I/K/L/M),第八轮 T 系列**待落地**(Batch N 建议见 §23.7);维持原判仅限设计决策项,见"十五、最终状态与遗留清单"与"十六、第六轮审计"。
|
||||
|
||||
---
|
||||
|
||||
@@ -974,3 +977,402 @@ Linux 基线(`aibox/venv`)维持 **490 passed** 不受影响。
|
||||
- `py_compile` 三文件通过
|
||||
|
||||
**Batch M 全部 5 个波次已收束**(S/F/W/M 四域,累计 §18~§22)。遗留"可不做/产品决策"项:F5(消息列表内存增长)、W9/W10(WS payload 上限/data 归一)、M6(并发保存 lost-update)、§17.5 表外 N 轮未完成项 N6(registry 转发总超时不完整)。
|
||||
|
||||
---
|
||||
|
||||
## 二十三、第八轮审计(2026-08-21,HEAD `3e860c1`)
|
||||
|
||||
方法:3 路并行深审(① `aibox.py` 编排层——并发/取消传播/orchestrator-worker 分发/子代理生命周期 ② `llm_client.py` + `model_config.py`——FailoverClient 状态机/流式解析/脱敏/key 回退 ③ `subsystems/` + `tools/` + `common/`——SQL/命令执行/路径/热重载/原子写/注册表并发)+ 主审对 v3.8 新代码(`process_manager.py` 731 行 + `process_tools.py` 408 行,commit 5cb3594)逐行复核 + Windows 测试基线确认。
|
||||
|
||||
### 23.1 Batch M 修复回归验证(全 5 波在位,零回归)
|
||||
|
||||
| 波次 | 核验结论 | 关键证据 |
|
||||
|---|---|---|
|
||||
| 波次 1 凭据收口 | 在位 | `skills/gitclient.py`/`ssh-exec.py`/`sftp-transfer.py`/`aes-crypto.py` 凭据改环境变量/凭据文件,无 `--password`/`--key` 命令行残留 |
|
||||
| 波次 2 路径边界 | 在位 | `gitclient.py:271-280` realpath 前缀 + `.git` 校验;`camera.py`/`facerec.py`/`vision.py` basename/扩展名校验在位 |
|
||||
| 波次 3 前端功能 | 在位 | `entry.html`/`frontend.html` 清空按钮先 cloneNode 再清空;三页面 manualDisconnect 清除重连定时器 |
|
||||
| 波次 4 命令面规整 | 在位 | `aibox.py` W1~W8/W11/W12 类型检查/钳制/缺参报错/空词提示/退出码/端口校验在位 |
|
||||
| 波次 5 客户端收尾 | 在位 | `voice_client.py` 单守护线程+Queue;`model_config.py` M2~M5/M7 规整;`client.py` finally disconnect |
|
||||
|
||||
**回归结论**:Batch M 全量零回归,4 个测试失败均为 Windows 平台环境项(symlink 权限 / python3 缺失 / POSIX 可执行位语义),与第七轮基线一致。
|
||||
|
||||
### 23.2 新代码首审——后台进程池 v3.8(commit 5cb3594)
|
||||
|
||||
v3.8 新增 `ProcessManager` 子系统 + 5 个 worker 执行面工具(`bg_run`/`bg_write`/`bg_list`/`bg_output`/`bg_kill`),弥补 `run_command` 同步阻塞、超时即杀的不足。`run_command` 原路径零改动。设计整体扎实(环形缓冲/三级终止/atexit 收尾/双闸开关/watermark 写读对齐均有测试锚定,39 项测试全过),但仍发现以下问题:
|
||||
|
||||
#### T8:`bg_write` watermark 取号与记录回显行非原子,并发写时写读对齐语义被破坏(高危)
|
||||
|
||||
`subsystems/process_manager.py:349-353`:
|
||||
|
||||
```python
|
||||
with entry.lock:
|
||||
watermark = entry.next_seq # ① 取号(持锁)
|
||||
# 锁释放后……
|
||||
echo = "[stdin] " + text.replace("\n", "⏎")
|
||||
self._record_line(entry, echo) # ② 记录回显行(再次取 entry.lock)
|
||||
```
|
||||
|
||||
两步之间释放了 `entry.lock`。两个并发 `bg_write`(A、B)的交错时序:A 取 `watermark=10` → B 取 `watermark=10` → A 记录回显 seq=10 → B 记录回显 seq=11。两个写入拿到相同 watermark,而 watermark 的核心语义是"`bg_output(since=watermark)` 必然且仅为本条输入的响应"——B 用 `since=10` 会同时读到 A 的响应,写读对齐被破坏。
|
||||
|
||||
虽然当前 worker 子代理在单线程 think 循环中串行调用工具、实际并发概率低,但 `ProcessManager` 注册表跨子代理全局共享(设计明确),且 `BgWriteTool` 无任何串行化保护,属正确性缺陷。**修法**:取号与记录回显行合并到同一临界区(在 `entry.lock` 内完成 `watermark = next_seq; next_seq += 1; 追加 echo 行`),约 5 行。
|
||||
|
||||
#### T28:`start()` 释放 self._lock 后 Popen 期间,concurrent kill/shutdown 找不到句柄但进程已启动(中危)
|
||||
|
||||
`subsystems/process_manager.py:139-229`:在 `self._lock` 内做完上限检查/分配 handle 后释放锁(:161),随后 `subprocess.Popen`(:168)可能耗时,直到 :219 才重新持锁把 entry 放入 `_procs`。此窗口内:
|
||||
- 另一线程 `shutdown()` 拿锁清空 `_procs`(:523)→ Popen 成功后 :221 检查 `_shutdown_done` 能兜底回收;
|
||||
- 但另一线程 `kill(handle=新handle)` 在 :567 `self._procs.get(handle)` 返回 None → 返回"句柄不存在",而进程实际已启动且无人管理,直到 reader EOF 或 atexit。
|
||||
|
||||
shutdown 分支有兜底,kill 竞态窗口窄(Popen 期间),定中危。**修法**:entry 在 Popen 前以 "starting" 状态放入 `_procs`,或 Popen 后再校验 `_shutdown_done` 时同时检查是否有并发 kill 意图(约 8 行)。
|
||||
|
||||
#### T22 关联:进程句柄表全局共享且无属主隔离(中危,见 T22)
|
||||
|
||||
#### T17:`text=True` 未指定 encoding,Windows cp936 解码 UTF-8 输出致 reader 线程静默退出(中危)
|
||||
|
||||
`subsystems/process_manager.py:178-180`:
|
||||
|
||||
```python
|
||||
text=True,
|
||||
bufsize=1,
|
||||
universal_newlines=True,
|
||||
```
|
||||
|
||||
Windows 上 text 模式默认按 locale(中文环境 cp936/GBK)解码。后台进程输出 UTF-8(Python `-u`、git、node 等)时 `readline` 抛 `UnicodeDecodeError` → `_reader` 的 `except Exception: pass`(:627-628)静默退出 → **该进程后续全部输出丢失**,bg_output 永远停在异常前的最后一行,且无任何错误提示。executors.py 有同型问题(第七轮未覆盖,本轮 T16)。**修法**:`encoding='utf-8', errors='replace'`(1 行)。
|
||||
|
||||
#### T18:`bg_run`/`bg_kill` 的 bool 参数未过 `_coerce_bool`,LLM 传字符串 `"false"` 恒真(中危)
|
||||
|
||||
`tools/process_tools.py:128-130, 386`:
|
||||
|
||||
```python
|
||||
def execute(self, command=None, ..., detach: bool = False, log_file: bool = False, ...):
|
||||
def execute(self, handle=None, force: bool = False, ...):
|
||||
```
|
||||
|
||||
`BgWriteTool` 对 `append_newline` 用了 `_coerce_bool`(:215),但 `BgRunTool.execute` 直接把 `detach`/`log_file` 透传给 manager(:136-137),`BgKillTool.execute` 直接透传 `force`(:396)。LLM 工具调用参数经 JSON 解析,布尔字段若传字符串 `"false"`(模型偶发)则 Python `bool("false") is True`——`detach` 误开启(aibox 退出不杀进程,进程脱离句柄管理)、`force` 误开启(跳过 SIGTERM 直接 SIGKILL,可能丢数据)、`log_file` 误开启。这正是第七轮波次 4(W 批命令面输入规整)要解决的问题类型,新工具漏了自家惯例。**修法**:三参数统一过 `_coerce_bool`(3 行)。
|
||||
|
||||
#### T21 关联:`env` 参数允许覆盖任意环境变量(中危,与 executors 同型,见 T21)
|
||||
|
||||
#### 低危——process_manager 专项
|
||||
|
||||
| # | 位置 | 问题 | 修法 |
|
||||
|---|---|---|---|
|
||||
| T-L1 | `process_manager.py:484-513` | `kill()` 主流程在 `self._lock` 外执行终止+轮询(最长 30s),与 `start()`/`shutdown()`/`list_procs()` 的 `_sweep_locked` 并发;`_sweep_locked` 在 kill 进行中可能将条目判为 running 并 TTL 终止(双终止),或 kill 完成后 status 被并发覆盖 | kill 期间持 entry 级状态锁,或 sweep 跳过正在 kill 的条目(~5 行) |
|
||||
| T-L2 | `process_manager.py:219-228` | shutdown 竞态分支 `self._terminate_tree(entry, force=True)` 后未关闭 `entry.log_fh`(正常路径在 reader finally 关闭,但 force kill 后 reader 线程可能未及时排空) | 该分支补 log_fh close(3 行) |
|
||||
| T-L3 | `process_manager.py:515-555` | `shutdown()` 无总预算:N 个进程 × (1s wait + 3s SIGTERM + 2s SIGKILL) 最坏 6N 秒,进程数达上限 16 时 atexit 收尾可阻塞 96 秒——atexit 期间若被 SIGKILL 则 detach 语义未记录 | 加总预算(如 15s),超时后 SIGKILL 全扫并记录(~8 行) |
|
||||
| T-L4 | `process_manager.py:373-382` | `bg_write` watchdog 线程在 write_timeout 后成为孤儿:超时返回后 watchdog 仍可能阻塞在 `stdin.write`(管道满 + 进程不读),daemon 线程随进程退出,但期间 `stdin_lock` 被持有,后续 write 全部超时 | 超时后无法安全杀死线程,文档已承认;可加"写入中毒"标志让后续 write 快速失败(~5 行) |
|
||||
| T-L5 | `process_manager.py:150-159` | `max_entries` 淘汰在 `self._lock` 内对 exited 列表排序(O(n log n)),n=32 影响极小,但与 running 上限检查的时序:`len(running) >= max_processes` 拒绝后未淘汰已退出条目即返回,已退出条目占满时无法启动新进程(即使 running 未达上限) | 上限检查前先 sweep+淘汰(:143 已 sweep,但淘汰在 :150 之后,调整顺序 2 行) |
|
||||
| T-L6 | `process_manager.py:182` | `cwd=os.path.expanduser(work_dir)` 未校验 work_dir 存在/可访问,不存在时 Popen 抛 FileNotFoundError 被 :187 捕获返回"命令未找到"——错误消息误导 | except 分支区分 cwd 不存在(~3 行) |
|
||||
|
||||
### 23.3 新发现——高危(11 项)
|
||||
|
||||
#### T1:WebSocket 服务无鉴权,`--host 0.0.0.0` 即暴露任意命令执行端点(设计层)
|
||||
|
||||
`aibox.py:3131` `websockets.serve(handler, host, port, ping_interval=None)` 无 token/来源校验;`:3220-3221` 支持 `--host/--port`;`:3520-3521` 覆盖配置。任何能建立 TCP 连接的客户端即获得完整智能体:经 `create_subagent` → worker → `run_command`(`shell=True`,工作目录默认 `~/.aibox/workspace` 但命令本身不限路径)= 任意命令执行。默认 localhost 尚可,`--host 0.0.0.0` 直接暴露到网络。read_file/write_file/list_files/web 均为"指引工具"(返回让模型调 run_command 的提示词),实际无路径白名单边界。**修法**(产品决策):握手 token(配置/CLI 生成,客户端首条消息校验);绑定非 localhost 时启动日志显著告警;若需多用户,为 worker 增加 workspace 沙箱路径限制。
|
||||
|
||||
#### T2:`set_proxy` 是编排者可调的全局持久化开关——LLM API Key 可被导向任意代理(设计层)
|
||||
|
||||
`aibox.py:1227` 编排者注册 `SetProxyTool`;`tools/proxy_tools.py:149` `save` 默认 True 持久化落盘;`llm_client.py:210-217` `proxies=self._resolve_proxies()` 让所有 LLM 请求走代理。`WORKER_ONLY_TOOLS` 拦截(`aibox.py:1416-1423`)只管住执行面工具,编排者合法持有 `set_proxy` 这一全局网络面:一次提示注入(网页内容/用户消息诱导)即可让 LLM 把代理指向攻击者服务器,此后**所有会话、定时任务、子代理的 LLM 请求都携带 `Authorization: Bearer <api_key>` 经过该代理**,且默认持久化跨重启生效。连通性验证只会验证代理可用——攻击者的代理当然可用。这是与 T22(bg_kill 全局无隔离)同型的"编排者合法持有但副作用全局"旁路,三层拦截(名字集合/未注册/拦截提示)共同管不住。**修法**(产品决策):`set_proxy` 移出 LLM 可直接调用的工具面(改需用户确认的命令);至少禁止持久化 save、代理变更时广播显式提示用户、对 api_url 域名加直连白名单。
|
||||
|
||||
#### T3:Windows 下 `shlex.quote` 指引无效且构成 cmd.exe 命令注入
|
||||
|
||||
`tools/basic_tools.py:109,183,257,339`:文件工具和 web 工具用 `shlex.quote`(POSIX 单引号语义)生成 `head/sed/tail/cat` 指引命令。当前开发/部署环境为 Windows,`run_command` 经 `subprocess.Popen(shell=True)` 由 **cmd.exe** 执行,cmd.exe 不认单引号,`& | ^ % < >` 仍是元字符:
|
||||
- 非恶意:URL 含 `&`(查询参数 `?a=1&b=2` 极常见)命令被截断、后半段被当命令执行报错——**功能必现故障**;`head/sed/tail` 等 POSIX 命令在 Windows 上根本不存在,四个文件工具+web 工具在当前环境整体不可用;
|
||||
- 恶意:网页内容提示注入诱导 LLM 抓取 `http://evil.com/?&del /f /s /q C:\data&`,指引命令夹带任意命令——命令注入成立(跨过 LLM 意图边界)。
|
||||
|
||||
**修法**:按 `os.name` 分支生成指引(Windows 走 PowerShell 等价命令并做 cmd 元字符转义),或更彻底地用 Python 直接实现文件读写/列目录(不再生成 shell 指引),从根上消除注入面。
|
||||
|
||||
#### T4:工具注册表共享输出队列 → 多会话并发时跨会话输出串扰(信息泄露)
|
||||
|
||||
`tools/registry.py:71-92`:
|
||||
|
||||
```python
|
||||
def set_output_queue(self, queue):
|
||||
self.current_output_queue = queue # 注册表级单值
|
||||
for tool in self._tools.values():
|
||||
tool.output_queue = queue # 覆写所有共享工具实例
|
||||
...
|
||||
tool.streaming_output = True
|
||||
tool.output_queue = self.current_output_queue # 并发调用互相覆盖
|
||||
```
|
||||
|
||||
`current_output_queue` 是单值,`output_queue` 是共享工具实例上的实例属性。多 WebSocket 连接/线程并发调用 `execute_tool` 时:会话 A 设置的队列被会话 B 覆盖 → **A 的工具输出推入 B 的队列(跨会话数据泄露)**;会话 B 结束清理(:166-167 无差别置 None)时把 A 正在使用的队列清空 → A 的输出静默丢失。PromptManager 已做 M6 原子替换、技能重载已做确定性处理,唯独会话级输出通路无并发防护。**修法**:队列随每次 `execute_tool` 调用局部传递(作为流式执行参数注入,而非实例属性);或注册表持有 `session_id -> queue` 映射,调用时显式指定并校验。
|
||||
|
||||
#### T5:共享 SubprocessCommandExecutor 单槽 `_current_process` → 并发执行时误杀/无法取消
|
||||
|
||||
`subsystems/executors.py:256-262, 299-305`:
|
||||
|
||||
```python
|
||||
with self._process_lock:
|
||||
self._current_process = process # 单槽位
|
||||
...
|
||||
finally:
|
||||
with self._process_lock:
|
||||
self._current_process = None # 先完成的线程清掉后来的引用
|
||||
```
|
||||
|
||||
`CommandLineTool` 全局共享一个 executor,但 `_current_process` 是单槽位。两个线程并发执行命令时:线程 2 的 Popen 覆盖线程 1 的引用 → 线程 1 的 `cancel()`/超时 `_terminate_tree` 会 **killpg 误杀线程 2 的整组进程**(POSIX 下连坐孙进程);线程 1 先结束的 finally 把引用清空 → 线程 2 的命令永远无法取消。`_process_lock` 只保护引用读写本身,没有保护"一个引用只属于一次执行"的语义。**修法**:每次执行生成调用令牌(自增 ID),`_current_process` 改 `{token: (proc, pgid)}`;`cancel(token)` 定向终止;finally 只删自己的 token。或对 executor 执行整体加锁串行化。
|
||||
|
||||
#### T6:FailoverClient 冷却键用 display_name,同名档案冷却状态互相污染
|
||||
|
||||
`llm_client.py:371-376` 冷却用 `client.name`(返回 `display_name or profile_name`,:639 display_name 优先);但 `create_llm_client` 的 P6 去重用 `profile_name`(:683)。display_name 无唯一性校验(model_config 只校验长度/可打印字符)。两个 profile_name 不同但 display_name 相同的档案(如主备两个 Kimi 端点都叫 "Kimi"):档案 A 失败进入冷却 → 档案 B 被 `_cooling` 误判为冷却中跳过——**该转移时不转移**。P6 注释表明作者意识到"冷却键同名冲突",但去重键与冷却键错位,防护完全无效。**修法**:冷却键改用 `id(client)` 或显式传入 profile_name,保证去重与冷却同一键。
|
||||
|
||||
#### T7:`FailoverClient.stream_chat` 是生成器,`_cancelled=False` 延迟到首迭代 → 取消标志丢失竞态
|
||||
|
||||
`llm_client.py:450-457`:`stream_chat` 含 `yield` 是生成器函数,函数体在首次 `next()` 才执行。时序:线程 A `gen = fc.stream_chat(msgs)`(仅创建生成器)→ 线程 B 调用 `fc.cancel()` 置 `_cancelled=True` → 线程 A 开始迭代 → :455 把 `_cancelled` 重置为 False——**用户的取消被清除,请求照常发出**(浪费 token)。这直接违反类文档(:336-339)声明的并发契约"cancel() 可从其他线程"。`_stats["requests"]` 计数同样因延迟执行而失真。**修法**:拆成普通方法 + 内部生成器(普通方法里先置 `_cancelled=False`/计数,再 return 内部生成器)。
|
||||
|
||||
#### T8:(见 §23.2,process_manager watermark 竞态)
|
||||
|
||||
#### T9:DeepSeekClient 非 200 路径响应不关闭,且 `_current_response` 残留
|
||||
|
||||
`llm_client.py:219-228`:
|
||||
|
||||
```python
|
||||
self._current_response = response # :219
|
||||
if response.status_code != 200: # :221 —— 在 try 之外
|
||||
if response.status_code == 429 or ...:
|
||||
raise TransientLLMError(...) # :224 —— 未 close、未清引用
|
||||
raise RuntimeError(...) # :227 —— 同上
|
||||
try: # :230 —— finally 只覆盖 200 路径
|
||||
...
|
||||
finally:
|
||||
self._current_response = None
|
||||
```
|
||||
|
||||
三层泄漏:(1) `stream=True` 的响应必须显式 close 或读完 body,否则底层 socket 不归还连接池;非 200 分支两个 raise 均未 close。(2) 异常的 `__traceback__` 持有 `stream_chat` 栈帧 → 栈帧持有局部 `response` → `FailoverClient` 捕获后 `last_exc = e`(:441/482)持续持有异常 → **整个转移链尝试期间(多个档案 × timeout 秒)所有失败响应的连接都不释放**。(3) `_current_response` 残留指向已失败响应,后续 `cancel()` 关陈旧引用。**修法**:把状态码检查移入 try/finally,非 200 分支 raise 前 `response.close()`。
|
||||
|
||||
#### T10:调度器 `await ai_future` 无硬超时兜底,卡死路径永久停摆调度循环
|
||||
|
||||
`aibox.py:319-365` `TaskScheduler.execute_task_as_conversation`:M3 修复(:351 `loop.call_later(task_timeout, _timeout_cancel)`)只做软取消——调用 `ai.cancel()`(Event + 关 HTTP 流 + kill 子进程)。若 AI 线程卡死在这三层都够不到的路径(Windows 代理连接黑洞、requests 读超时被服务器滴流输出规避、kill 失败的进程),`run_ai` 永不回调 `set_result`,`await ai_future`(:353)**永久挂起** → `TaskScheduler.run` 的 while 循环卡死 → 所有后续定时/即时任务停摆,且任务已 `mark_as_triggered` 不会重试。对比 `--once` 模式已有 `_hard_exit`(`os._exit(4)`,:3441-3449)硬杀兜底,调度器恰恰遗漏。**修法**:`response = await asyncio.wait_for(ai_future, timeout=task_timeout + 30)`;超时后告警并继续循环;可参照 --once 增加二级 Timer 硬兜底。
|
||||
|
||||
#### T11:Windows 进程树收割缺失(executors + process_manager 两处同型)
|
||||
|
||||
`subsystems/executors.py:85-97` 和 `subsystems/process_manager.py:600-618`:POSIX 用 `start_new_session` + `killpg` 收割整棵进程树(L9 修复),但 Windows 分支只对单个包装进程 `terminate()/kill()`。Windows 上 `shell=True` 生成 `cmd.exe /c`,`start xxx`、PowerShell 后台任务、dev server spawn 的子进程等孙进程不会被终止:超时/取消/kill 后孙进程继续运行成孤儿,持续占用 CPU/文件锁/端口。当前环境即 Windows,此缺陷为实际生效状态;process_manager 管理的恰是会 spawn 子进程的长驻服务(dev server/训练任务),影响更甚。**修法**:Windows 上用 Job Object(子进程整个生命周期绑定作业对象,作业销毁即收割全树),或至少超时/终止路径追加 `taskkill /PID {pid} /T /F`。
|
||||
|
||||
### 23.4 新发现——中危(17 项,T12~T28)
|
||||
|
||||
#### T12:`has_stream_output` 恒为 False(死逻辑):流式输出重复推送 + 超长结果绕过 N8 预算直发前端
|
||||
|
||||
`aibox.py:1427-1462`::1429 排空 output_queue("complete" 标记在此被消费丢弃)→ :1451 `has_stream_output = not output_queue.empty()` **排空之后判定,恒为 False**。后果:(1) 流式工具(run_command 等)的完整过程输出(line 流)与最终 result(含完整"输出"字段)**重复展示两遍**;(2) 超长工具结果在 :1458 丢弃**之前**已在 :1454 `output_callback("line", result + "\n")` 全量推送给前端——registry 流式通道有 N8 预算限制(`registry.py:120-147`),这条非流式直发路径完全绕过输出限制。**修法**:在 :1429 排空循环中记录 `saw_complete = (msg_type == "complete")`,用其替代 `not output_queue.empty()`;长度截断移到 output_callback 发送之前。
|
||||
|
||||
#### T13:异常消息内嵌未脱敏 `response.text`,上抛路径无统一脱敏
|
||||
|
||||
`llm_client.py:224-228` 构造异常时直接内嵌 `response.text[:200]`,:490-491 `raise last_exc` 原样上抛。代码已承认威胁模型(:29-31 L12 注释:网关/WAF 错误页会回显 Authorization 头),但 `_sanitize_log_msg` 只应用在 `_log`、`on_failover` 回调、model_config 三处。任何调用方 `except Exception as e: logger.error(str(e))` 未自觉脱敏,令牌即落盘。**脱敏应在异常构造点一次完成**,而非依赖每个下游消费方。**修法**:构造异常时统一 `_sanitize_log_msg(response.text[:200])`。
|
||||
|
||||
#### T14:FTS5 MATCH 左操作数疑似用错表名,FTS5 全程静默降级走 LIKE(待运行验证)
|
||||
|
||||
`subsystems/db.py:188-195`:
|
||||
|
||||
```python
|
||||
FROM memories m
|
||||
JOIN memories_fts f ON m.rowid = f.rowid
|
||||
WHERE memories_fts MATCH ? -- 表已别名 f,LHS 用原表名
|
||||
```
|
||||
|
||||
FTS5 `MATCH` 左操作数必须是 FROM 子句中的表名或别名。`memories_fts` 已被别名 `f` 遮蔽,`memories_fts MATCH ?` 极可能抛 `no such column: memories_fts`,被 :196-198 的 `except OperationalError` 静默吞掉 → **FTS5 索引与 rank 排序从未生效,全部检索走 LIKE 兜底**,且每次搜索刷一条 warning。此项需运行验证(SQLite 版本可能对表名 MATCH 有宽松容忍)。另:空关键词 `""` 生成空短语 → FTS 语法错 → LIKE `%%` 全表匹配。**修法**:改 `WHERE f MATCH ?`;入口过滤空关键词。建议先写复现用例确认。
|
||||
|
||||
#### T15:非流式工具结果完全不受 TOOL_MAX_RESULT_LENGTH 约束(输出限制绕过)
|
||||
|
||||
`tools/registry.py:173-175` 普通执行分支 `result = tool.execute(**kwargs)` 直接返回,无任何长度截断。N8 修复只给流式转发加了 `line_budget`(:120-147)。无输出队列的场景(后台任务、非流式会话)中,`memory get`/`memo get` 返回不设上限的 `content`(`add_memory`/`add_memo` 入库时 content 长度也无校验)→ 超长内容直接灌入对话上下文,token/内存放大。流式与非流式两条路径限制不一致即绕过。**修法**:普通执行分支同样套用 `TOOL_MAX_RESULT_LENGTH` 截断;入库时加 content 长度上限。
|
||||
|
||||
#### T16:executors `text=True` 未指定 encoding,Windows cp936 解码 UTF-8 致输出丢失(与 T17 同型)
|
||||
|
||||
`subsystems/executors.py:240-254, 333-347`:`text=True` 未指定 encoding,Windows 中文环境按 cp936 解码。子进程输出 UTF-8 时 readline 抛 `UnicodeDecodeError` → reader 线程 `except Exception: pass`(:135-137)静默退出 → **命令输出部分或全部丢失**,GBK 恰可解码的部分则乱码。**修法**:显式 `encoding='utf-8', errors='replace'`。
|
||||
|
||||
#### T17:(见 §23.2,process_manager encoding 同型)
|
||||
|
||||
#### T18:(见 §23.2,process_tools bool 参数未规整)
|
||||
|
||||
#### T19:`RequestException` 捕获过宽,永久配置错误被当瞬时故障转移+冷却
|
||||
|
||||
`llm_client.py:437, 477` `except (TransientLLMError, requests.exceptions.RequestException)` 捕获全部 requests 异常基类。档案 `api_url` 手写错(漏 scheme)时 `requests.post` 抛 `MissingSchema`——永久性错误被当瞬时:标记冷却 60s、触发转移、计入 failovers。每次冷却到期都会再撞一次坏端点,统计与冷却表被永久错误污染。**修法**:显式枚举瞬时异常 `(ConnectionError, Timeout, ChunkedEncodingError)`。
|
||||
|
||||
#### T20:不完整流被静默视为成功,回答无声截断无信号
|
||||
|
||||
`llm_client.py:238-245, 476`:代理缓冲后提前 FIN / 服务端中途断开时,末行半截 JSON 解析失败 `continue`(:244),迭代耗尽 → FailoverClient :476 将"从未产出 `[DONE]` 的自然结束"视为成功。业务层收到**无声截断的回答**,无任何错误信号,回复内容静默丢失尾部。**修法**:跟踪 `saw_done`,正常结束但未收到 `[DONE]` 时抛 `TransientLLMError`(首事件前可转移,已产出事件则上抛)或 yield error 事件。
|
||||
|
||||
#### T21:`run_command`/`bg_run` 的 env 参数允许覆盖任意环境变量(注入放大面)
|
||||
|
||||
`subsystems/executors.py:234-236, 327-329` 和 `subsystems/process_manager.py:163-165`:`exec_env = os.environ.copy(); exec_env.update(env)`,env 在工具 schema 中是开放 object,LLM(可被网页内容提示注入诱导)可传 `{"LD_PRELOAD": "...", "PATH": "...", "PYTHONPATH": "..."}`,在 `shell=True` 之上进一步扩大代码执行/劫持面(POSIX 下 LD_PRELOAD 可加载任意共享库)。**修法**:env 键白名单,或至少黑名单拒绝 `LD_*`/`DYLD_*`/`PATH`/`PYTHONPATH`/`SHELL`/`PSModulePath`;增加审计日志。
|
||||
|
||||
#### T22:定时任务 AI 以 orchestrator profile 运行且共享全局 PromptManager/DB,无人值守时段可改写人格/代理/记忆
|
||||
|
||||
`aibox.py:104-125` `_create_task_ai` 未传 `tool_profile`,默认 `TOOL_PROFILE_ORCHESTRATOR`(:1027),且与主对话**共享** `prompt_manager`/`memory_db`/`memo_db`,注册了 `UpdateIdentityTool`/`UpdateSoulTool`(改写全局系统提示词)、`set_proxy`(见 T2)、`MemoTool`。任务内容由备忘录标题/正文驱动(可能源自 LLM 或被网页内容污染的生成物),focus_prompt(:274-284)只是软约束——被污染的任务可在无人值守时静默改写人格提示词、代理与记忆库。上下文隔离做了(独立 messages/executor),但**权限未隔离**。关联:进程句柄表全局无属主隔离(aibox.py:1185-1190 注释自述"跨子代理可管理彼此启动的进程"),任意 worker 可 `bg_kill` 其他会话进程。**修法**:任务 AI 显式传精简 `tool_profile`(无 update_identity/update_soul/set_proxy),或传入 PromptManager 只读快照;ProcessManager 条目增加 owner 隔离。
|
||||
|
||||
#### T23:对话压缩可产生非法消息序列(assistant.tool_calls 与 tool 结果被切开)
|
||||
|
||||
`subsystems/conversation_compressor.py:100-161, 226-228, 370-394`:`_is_range_safe` 恒返回 True(不保护 tool 消息),压缩块边界(按条数切)可能落在带 `tool_calls` 的 assistant 消息与其 tool 结果消息之间。块被摘要替换后,保留的 assistant 消息的 `tool_calls` 找不到配对的 tool 消息——OpenAI 协议要求二者成对,**后续 API 调用直接 400 报错**。`_validate_message_structure` 只修复了反向(孤立 tool 消息)。**修法**:压缩块边界对齐完整对话轮次(user → assistant → tool* 整体进出块);或压缩后剥离失去配对的 `assistant.tool_calls` 字段。
|
||||
|
||||
#### T24:非流式路径下 `timeout<=0` 可无限阻塞,registry 3600s 兜底不生效
|
||||
|
||||
`subsystems/executors.py:147-149` 注释自认 timeout<=0 语义为"不设超时",并称"registry 层 3600s 流式兜底仍生效",但 registry 的兜底只在流式分支(`registry.py:114-127`);非流式分支(:173-175)`result = tool.execute(**kwargs)` 无任何超时兜底。LLM 传 `timeout=0` 时,一条挂起的命令(`python -c "import time; time.sleep(1e9)"`)将**永久阻塞调用线程**。**修法**:非流式分支同样用线程+deadline 包裹;或禁用 timeout<=0 语义(下限夹取到 1s),由 registry 统一兜底。
|
||||
|
||||
#### T25:MemoryDatabase 漏设 busy_timeout,与 MemoDatabase 共库时写锁等待不一致
|
||||
|
||||
`subsystems/db.py:38, 121, 142, 169, 243, 259, 277`:MemoryDatabase 全部路径直连 `sqlite3.connect(self.db_path)`,无 `busy_timeout`;而 MemoDatabase._connect(:307-312)已统一 `PRAGMA busy_timeout = 10000`(J9 修复)。两个类默认共用同一个 `memories.db`,J9 只覆盖了 MemoDatabase。CLI `--memo-add` 与服务进程并发写时,MemoryDatabase 路径只靠 sqlite3 默认 5 秒等待,锁冲突即抛 `database is locked`,记忆保存/检索失败。**修法**:MemoryDatabase 同样走 `_connect()` 工厂(busy_timeout + WAL 显式设置)。
|
||||
|
||||
#### T26:SkillManager.reload 非原子替换,热重载期间读者看到空/半截技能表
|
||||
|
||||
`subsystems/skill_manager.py:127-128`:`_load_skills` 先 `self.skills = {}` 清空再逐个重建。与 PromptManager M6 修复(构建新列表后原子替换)完全同型的问题在这里未修。reload 期间(每个可执行技能 `--help` 最多 10 秒,技能多时窗口很长),并发线程调用 `search_skills`/`get_skill_content` 会读到空字典或半截数据——技能"消失"再"逐渐出现"。**修法**:本地构建 `new_skills` 完成后一次性 `self.skills = new_skills`。
|
||||
|
||||
#### T27:atomic.py 在 Windows 上 os.replace 目标被占用即失败(高频写路径必踩)
|
||||
|
||||
`common/atomic.py:26` `os.replace(tmp_path, path)`:POSIX 下无条件成功;Windows 下若目标文件正被其他进程打开(杀毒扫描、编辑器、另一进程正在读),`MoveFileExW(REPLACE_EXISTING)` 抛 `PermissionError`。`conversation_history.add_message` **每条消息**都原子写 `current_conversation.json`(写频极高),与扫描器/备份工具冲突概率不低——一次冲突即该条对话持久化失败。**修法**:捕获 `PermissionError` 做短退避重试(3 次 × 50ms),仍失败再上抛;写侧考虑降频(批量/定时快照)。
|
||||
|
||||
### 23.5 新发现——低危与信息(精选 20 项)
|
||||
|
||||
| # | 位置 | 问题 | 修法 |
|
||||
|---|---|---|---|
|
||||
| T-L7 | `aibox.py:1469-1474, 383-384, 2902-2903` | 多处静默吞异常(`add_tool_call` 失败无日志;裸 `except: pass` 连 KeyboardInterrupt 一起吞) | `except Exception as e: logger.warning(...)`,禁裸 except |
|
||||
| T-L8 | `aibox.py:1061, 1857, 1866, 1883` | `stats["errors"]` 列表无界增长,长生命周期会话内存缓慢增长 | `collections.deque(maxlen=50)` |
|
||||
| T-L9 | `aibox.py:1813` | 主对话 orchestrator 工具循环无轮次上限,每轮递归加深调用栈;子代理有 max_turns 钳制,主对话没有——300+ 轮触发 RecursionError(虽被 catch 回滚但烧钱) | 加软轮次上限(如 50),超限注入"请总结并结束" |
|
||||
| T-L10 | `aibox.py:3114, 3121, 1239` | 启动期技能扫描重复执行 2-3 遍(ws_handler/scheduler/add_ai 各自 SkillManager),每个可执行技能串行 --help,技能多时启动延迟放大 | EnhancedAIManager 持有单例注入三处 |
|
||||
| T-L11 | `aibox.py:362-365` | 调度器 AI 线程 join(timeout=1.0) 后成孤儿,任务"超时完成"后 output_callback 仍 run_coroutine_threadsafe 广播迟到输出(前端 complete 后又收到 chunk) | 回检查 `timed_out.is_set()` 丢弃 |
|
||||
| T-L12 | `aibox.py:1302-1336` | `_save_tool_result_to_file` 已无调用方(超长改为直接丢弃),死代码——本可用于"超长落盘+返回路径"的更好体验(与 T12 相关) | 删除,或反向利用落盘返回路径 |
|
||||
| T-L13 | `aibox.py:2130-2150` 等 | DeepSeekChat 实例(会话/任务/子代理)无 dispose 钩子,LLM client/executor/registry 全靠 GC | 加 `close()`(cancel + queue 置 None + executor 清理),unregister/任务 finally 调用 |
|
||||
| T-L14 | `llm_client.py:235-237` | SSE 解析只认 `"data: "`(带一个空格),SSE 规范允许 `data:foo` 无空格;非主流网关/本地 vLLM 可能发无空格格式致整条流静默丢弃 | `data = line[5:]; if data.startswith(" "): data = data[1:]` |
|
||||
| T-L15 | `llm_client.py:536-538` | `resolve_api_key` key 文件回退 `except Exception` 吞原始异常,不区分文件不存在/权限拒绝,丢失原异常链 | `raise RuntimeError(...) from e`,消息区分 key_path 是否存在 |
|
||||
| T-L16 | `llm_client.py:246-251` | `finish_reason` 被完全忽略,业务层无法区分"答完了"与"被 max_tokens 截断了" | done 事件携带 finish_reason |
|
||||
| T-L17 | `model_config.py:317-322` | 添加非 DeepSeek 档案时环境变量名默认 `DEEPSEEK_API_KEY`,一路回车会把 DeepSeek 凭据发给其他厂商(moonshot/aliyuncs);正是 resolve_api_key 文档声称要杜绝的跨厂商凭据外发 | 默认值留空强制输入,或按 api_url 域名推荐 |
|
||||
| T-L18 | `model_config.py:289-295` | 展示层 `list_profiles` 原样 print 档案名/display_name/api_url,手写 config 嵌 ANSI 转义(`\x1b]0;pwned\x07`)每次列表即触发;写入侧有 isprintable 校验,读出侧裸奔 | 输出前过滤控制字符 |
|
||||
| T-L19 | `subsystems/skill_manager.py:70-71` | Markdown 技能文件 `f.read()` 全量读取无大小限制(对比 `_run_help` 有 2000 字符截断);skills 目录可经 run_command 写入,数 GB .md 全量入内存+进系统提示词 | 读取前 getsize 检查,超限拒绝加载并告警 |
|
||||
| T-L20 | `common/constants.py:220-221` vs `config.py:165` | `TOOL_MAX_RESULT_LENGTH` fallback 100000 而 config 默认 10000——用户删掉配置键时上限放大 10 倍;`MAX_HISTORY_CONTEXT` fallback 50 vs 默认 120 同型 | 统一 fallback 与 config 默认值 |
|
||||
|
||||
**其余低危**(简列):`common/atomic.py` 未 fsync 父目录 + mkstemp 0600 权限未恢复(POSIX);`tools/registry.py` 注册/迭代无锁(启动期注册风险低);`conversation_history.py` conv_id 秒级时间戳可重 + 归档 TOCTOU + 每条消息全量重写(O(n²));`prompt_manager.py` prompts 目录直读无大小限制;`skill_manager.py` Windows 把 .py/.sh 判可执行但 CreateProcess 无法直接执行(WinError 193 被吞);`skill_tools.py` InstallSkill 目标文件名无字符校验(对比 CreateSkill 有正则);`db_tools.py` 空关键词 LIKE `%%` 全表匹配;`logger.py` 初始化竞态 + 多进程 RotatingFileHandler 轮转冲突;`config.py` set_proxy/_save_config 无锁同进程竞态 + 跨进程 lost-update;`executors.py` `except subprocess.TimeoutExpired` 不可达(Popen 无 timeout);`conversation_compressor.py` `weight_decay_factor` 硬编码 0.6 致 constants 成死配置;`base.py` `check_path_safety` 名不副实(只查文件大小)。
|
||||
|
||||
### 23.6 查过但无问题(摘要)
|
||||
|
||||
- **v3.8 后台进程池**:环形缓冲 seq 单调连续性(淘汰最旧行仍保持 seq 连续,since 游标续读正确);三级优雅终止(关 stdin→SIGTERM 宽限→SIGKILL)逻辑正确;atexit 收尾幂等;detach 语义正确(shutdown 保留运行仅注销);双闸开关(commands.enabled + background_processes.enabled);单行截断 16KB + 单次拉取 1MB 硬顶;watermark wait_lines expect 式同步等待 + 进程退出提前返回;写入 watchdog 防管道满阻塞;39 项测试覆盖生命周期/游标续读/环形淘汰/watermark 对齐/EOF 优雅终止/log tee/工具分发,全过
|
||||
- **llm_client**:免密端点密钥隔离(api_key_env 空串不读 key 文件);用户中断绝不转移(_cancelled 标志,P 批修复);fallback 链硬性校验(指向不存在档案抛错、缺 model 跳过告警);Bearer 日志脱敏在位(L12);Retry-After cap 3600s(P2);fallback 去重(P6)
|
||||
- **aibox.py 编排**:N2 drain 线程批次级生命周期;N14 任务集合登记;R3 注销回收闭环;J7 广播超时;消息配对自愈(`_repair_tool_pairs`/`_validate_enhanced_messages`);失败回滚(`_rollback_messages`);三层取消的第一层(Event)检查点完备(流循环/工具调用前/think 前置);N5 共享父取消事件不再在子代理首轮 clear(H1 回归修复在位)
|
||||
- **subsystems**:db SQL 全参数化 + FTS5 引号翻倍 + LIKE `%/_/\` 转义(W4/第七轮);MemoDatabase J9 busy_timeout + WAL;atomic 写单次原子性;PromptManager M6 原子替换;技能路径 realpath 边界(J12/L4)
|
||||
|
||||
### 23.7 修复建议(Batch N,按优先级分波)
|
||||
|
||||
| 波次 | 内容 | 项 | 预估规模 |
|
||||
|---|---|---|---|
|
||||
| 波次 1 速赢 | 几行改动高收益:非 200 响应 close(T9)、异常构造点脱敏(T13)、bool 参数规整(T18)、encoding=utf-8(T16/T17)、saw_complete+输出预算(T12)、watermark 原子取号(T8) | T8/T9/T12/T13/T16/T17/T18 | ~30 行 |
|
||||
| 波次 2 并发正确性 | registry 输出队列调用维度隔离(T4)、executor 进程槽令牌化(T5)、FailoverClient 冷却键(T6)、生成器拆分(T7)、非流式超时兜底(T24)、非流式输出限制(T15) | T4/T5/T6/T7/T15/T24 | ~80 行 |
|
||||
| 波次 3 平台兼容 | Windows 进程树 taskkill/Job Object(T11)、shlex 指引跨平台(T3)、atomic Windows 重试(T27)、MemoryDatabase busy_timeout(T25)、SkillManager 原子 reload(T26) | T3/T11/T25/T26/T27 | ~70 行 |
|
||||
| 波次 4 LLM 健壮性 | RequestException 收窄(T19)、不完整流 saw_done(T20)、调度器硬超时(T10)、env 键白名单(T21)、压缩块边界对齐(T23)、FTS5 别名修复+空词(T14,先写复现用例) | T10/T14/T19/T20/T21/T23 | ~60 行 |
|
||||
| 波次 5 权限隔离(产品决策) | WS 握手 token(T1)、set_proxy 移出工具面/禁持久化(T2)、定时任务精简 profile + ProcessManager owner 隔离(T22)、model_config env 默认值/ANSI 过滤(T-L17/T-L18) | T1/T2/T22 + 低危 | ~80 行 |
|
||||
|
||||
**定级说明**:
|
||||
- T1/T2/T10 为设计层问题,需产品决策(是否支持多用户/网络暴露、是否允许 LLM 自改代理、调度任务权限边界),不阻塞波次 1~4
|
||||
- T3 在 Windows 开发机上是**必现功能故障**(文件/web 工具指引命令不可用),不是纯理论问题,建议优先
|
||||
- T4/T5 是多连接 WebSocket 服务的正确性前提,虽然当前单用户场景触发概率低,但属于"架构假设与部署形态不匹配"
|
||||
- T14(FTS5 失效)必须先运行复现确认,若属实则是长期静默降级(记忆搜索质量下降但不报错)
|
||||
- T8 watermark 竞态在当前串行 think 下概率低,但 ProcessManager 全局共享的设计声明使其成为正确性债务
|
||||
|
||||
### 23.8 测试基线(Windows 开发机)
|
||||
|
||||
审计时全量 **532 passed / 4 failed / 1 skipped**(40.9s),较第七轮(493 passed)增加 39 项(v3.8 `test_process_manager.py`)。4 失败均为 Windows 平台环境项:
|
||||
|
||||
1. `test_h_fixes.py::TestL4PersonaSymlink::test_symlink_outside_skills_dir_rejected` — Windows symlink 需管理员权限(OSError WinError 87)
|
||||
2. `test_h_fixes.py::TestL5NoPartialChar::test_truncated_output_has_no_replacement_char` — Windows 无 `python3` 命令
|
||||
3. `test_k_fixes.py::TestJ16SkillManager::test_markdown_skipped_when_executable_same_name` — POSIX 可执行位语义
|
||||
4. `test_k_fixes.py::TestJ16SkillManager::test_help_output_truncated` — 同上
|
||||
|
||||
与第七轮基线完全一致,零回归。
|
||||
|
||||
### 23.9 第八轮结论
|
||||
|
||||
3 路并行深审 + 主审对 v3.8 新代码逐行复核:**Batch M 全 5 波零回归**;后台进程池 v3.8 新代码设计扎实、测试充分(39 项全过),但仍发现 watermark 竞态(T8)、encoding 缺失(T17)、bool 参数未规整(T18)等问题。全模块新发现 **T1~T28 中高危 28 项 + 低危与信息 20+**,其中高危 11 项集中于三个系统性主题:
|
||||
|
||||
1. **平台矩阵未声明**:大量 POSIX 专用逻辑(shlex 指引/killpg/heredoc/无 encoding/atomic replace)在 Windows 部署下要么失效(T3 文件工具必现故障、T16 输出丢失)要么成为注入面(T3 cmd.exe 注入),要么功能受限(T11 孙进程残留、T27 原子写失败)。建议要么全面适配 Windows,要么在文档/启动检查中声明仅支持 POSIX。
|
||||
2. **"单会话串行"隐含假设在多连接 WebSocket 下不成立**:registry 输出队列(T4)、executor 进程槽(T5)、watermark(T8)都假设同一时刻只有一个会话在执行工具,但 WebSocket 服务天然多连接并发。需按调用维度(token/session_id)隔离。
|
||||
3. **防线位置错位**:脱敏建在消费侧而非产生侧(T13)、冷却键与去重键不一致(T6)、校验建在写入侧而非读出侧(T-L18)、生成器延迟副作用破坏取消契约(T7)、`has_stream_output` 排空后判定恒 False(T12)。
|
||||
|
||||
八轮累计发现 **166+ 项**(高危 18 / 中危 54 / 低危与信息 94+);前七轮 118 项全部修复落地(Batch E/F/G/H/I/K/L/M),第八轮 T 系列待落地(Batch N 建议见 §23.7)。
|
||||
|
||||
## 二十四、Batch N 修复记录(波次 1~5,2026-08-21)
|
||||
|
||||
范围:§23.7 全部五个波次——速赢 7 项 + 并发正确性 6 项 + 平台兼容 5 项 + LLM 健壮性 5 项修复(另有 T14 复现证伪 1 项)+ 权限隔离 3 项修复与 2 项产品决策收束。
|
||||
|
||||
### 24.1 波次 1 落地(速赢批)
|
||||
|
||||
| # | 修复内容 | 文件/位置 |
|
||||
|---|---|---|
|
||||
| T8 | watermark 取号与回显行记账移入同一临界区(`entry.lock` 改 `RLock` 防重入死锁),保证 watermark 等于回显行真实序号 | `subsystems/process_manager.py:79,348-354` |
|
||||
| T9 | 流式/非流式两条非 200 路径抛出前 `response.close()`,流式路径同时清 `_current_response`,消除连接占用与残留引用 | `llm_client.py:172-190,233-247` |
|
||||
| T13 | `_sanitize_log_msg` 增加裸 `sk-xxx` 正则;两处 `response.text[:200]` 入异常前先脱敏 | `llm_client.py:33-42` |
|
||||
| T12 | drain 线程与后排空循环共享 `saw_complete` 标志替代排空后恒 False 的 `empty()` 判定;长度截断移到 output_callback 之前,堵住非流式直发绕过 N8 预算的通道 | `aibox.py:1379-1395,1432-1475` |
|
||||
| T16 | executors 两处 Popen 显式 `encoding="utf-8", errors="replace"` | `subsystems/executors.py` |
|
||||
| T17 | process_manager Popen 同款显式编码 | `subsystems/process_manager.py:181-182` |
|
||||
| T18 | `detach`/`log_file`/`force` 三参数统一过 `_coerce_bool` | `tools/process_tools.py:136-139,401-402` |
|
||||
|
||||
### 24.2 波次 2 落地(并发正确性批)
|
||||
|
||||
| # | 修复内容 | 文件/位置 |
|
||||
|---|---|---|
|
||||
| T4 | `set_output_queue` 不再把队列传播到共享工具实例(跨会话串扰的覆盖向量,`execute_streaming` 本就每次注入局部队列);`execute_tool` 入口快照会话队列,转发循环全程用快照——并发会话互不覆盖、结束清理不影响在途转发 | `tools/registry.py:82-92,139-157` |
|
||||
| T5 | executor 单槽 `_current_process` 改令牌多槽注册表 `{token: (proc, pgid)}`;`_terminate_tree` 改显式传 `(proc, pgid)` 不再读共享 `_current_pgid`(误杀根因);finally 只注销自己的槽位;`cancel()` 语义为终止本执行器全部在途命令(原为"杀单槽里的随机进程") | `subsystems/executors.py:72-153,286-351,377-459` |
|
||||
| T6 | FailoverClient 冷却键改 `id(client)`(`_cool_key`),与 P6 去重键语义一致——同名 display_name 档案不再互相污染冷却态;`test_models_profile.py` 5 处种键方式同步更新 | `llm_client.py:390-410`;`tests/test_models_profile.py:823-825,1001-1003,1014,1106,1235-1242` |
|
||||
| T7 | `stream_chat` 拆为普通方法(同步完成 `_cancelled=False` 重置 + 计数)+ 内部生成器 `_stream_chat_impl`——创建生成器与首迭代之间的 `cancel()` 不再被延迟执行的重置清掉 | `llm_client.py:478-501` |
|
||||
| T15 | registry 两条分支的最终返回值统一过 `_truncate_result`(TOOL_MAX_RESULT_LENGTH),非流式结果不再无上限直入上下文,与 N8 line_budget 同口径 | `tools/registry.py:71-80,190-193` |
|
||||
| T24 | 非流式 `execute()` 的 timeout 下限夹取到 1s(含字符串容错)——该路径无 registry 3600s 兜底,timeout<=0 会永久阻塞调用线程;流式路径保留 0=不限语义 | `subsystems/executors.py:267-273` |
|
||||
|
||||
### 24.3 测试基线(Windows 开发机)
|
||||
|
||||
波次 1 后:**532 passed / 4 failed / 1 skipped**(43.8s);波次 2 后:**532 passed / 4 failed / 1 skipped**(41.5s)。4 失败均为 §23.8 所列同一批 Windows 平台环境项,零回归。
|
||||
|
||||
### 24.4 待落地
|
||||
|
||||
无——波次 1~5 全部落地。T1/T2 的产品决策记录见 §24.9。
|
||||
|
||||
### 24.5 波次 3 落地(平台兼容批,2026-08-21)
|
||||
|
||||
范围:T3/T11/T25/T26/T27,共 5 项。
|
||||
|
||||
| # | 修复内容 | 文件/位置 |
|
||||
|---|---|---|
|
||||
| T3 | `basic_tools.py` 新增平台感知引用 `_quote_arg()`(POSIX=shlex.quote;Windows=双引号包裹 + 拒绝 `"`/换行/`%变量名%` 形态——cmd 变量展开是经典注入向量,URL 编码 %20 与 PowerShell `%{http_code}` 不匹配该模式不受影响)与 `_ps_quote()`(PowerShell 单引号字面量);四个工具按 `os.name` 分支生成指引:ReadFile 走 Get-Content/Select-Object(head/tail/行范围/wc 等价),WriteFile 走 Base64+Set-Content/Add-Content(heredoc 在 cmd 下不可用:命令行不能含换行;Base64 字符集不含任何 cmd/PowerShell 元字符,多行/引号/特殊字符全覆盖;超 6000 字符提示分块追加),ListFiles 走 Get-ChildItem,Web 走系统内置 curl.exe(分析类动作 text 给 PowerShell 管道等价,title/links/search 指引取原始 HTML 自行分析);`_proxy_curl_args`/search keyword 引用同步切换 | `tools/basic_tools.py:33-78,150-237,297-343,389-414,474-533` |
|
||||
| T11 | executors 与 process_manager 两处 `_terminate_tree` Windows 分支改 `taskkill /PID {pid} /T /F` 递归终止整棵进程树(无 /F 仅对 GUI 进程有效,控制台进程杀不掉,故统一强制;force 参数在 Windows 不区分),taskkill 失败(进程已退出/无权限)兜底单进程 kill;POSIX 分支行为不变 | `subsystems/executors.py:109-157`;`subsystems/process_manager.py:604-634` |
|
||||
| T25 | MemoryDatabase 新增 `_connect()` 统一连接工厂(`PRAGMA busy_timeout = 10000`),7 处直连 `sqlite3.connect` 全部换走工厂——与 MemoDatabase J9 修复对齐,两类共库时锁等待策略一致,CLI `--memo-add` 与服务进程并发写不再抛 `database is locked` | `subsystems/db.py:36-45`(工厂)+ 7 处调用点 |
|
||||
| T26 | `SkillManager._load_skills` 本地构建 `new_skills` 完成后一次性原子替换 `self.skills`(对齐 PromptManager M6)——reload 期间(每个可执行技能 `--help` 最多 10s)并发读者不再看到空/半截技能表;目录不存在分支同样保持显式赋值语义 | `subsystems/skill_manager.py:125-213` |
|
||||
| T27 | `atomic_write_text` 的 `os.replace` 捕获 `PermissionError` 做 3 次短退避重试(50/100ms)——Windows 下目标被杀毒扫描/编辑器/并发读者占用时一次冲突不再直接丢该条持久化(conversation_history 每条消息都走此路径);重试耗尽上抛原异常,临时文件清理逻辑不变;POSIX 下 `os.replace` 无条件成功不受影响 | `common/atomic.py:15-52` |
|
||||
|
||||
**Windows 分支验证方式注记**:Linux 环境以 `mock.patch("os.name", "nt")` 模拟验证命令生成逻辑(引用规则/拒绝路径/命令形态),生成的 PowerShell/curl.exe 命令未在真实 Windows 上执行验证,Windows 开发机建议抽测指引命令的实际可用性。
|
||||
|
||||
### 24.6 测试基线(Linux,`aibox/venv`)
|
||||
|
||||
波次 3 后:**566 passed / 1 failed**(30.1s)。新增 `tests/test_n_fixes.py` 30 项(T3 引用规则 6 + Windows 指引生成 10 + T11 taskkill 4 + T25 busy_timeout 2 + T26 原子 reload 2 + T27 重试 3 + POSIX 锚定 3);1 失败为 §11.3 所列 Windows 向用例(`test_execute_timeout`,本机无 ping 命令),与代码无关,零回归。
|
||||
|
||||
### 24.7 波次 4 落地(LLM 健壮性批,2026-08-21)
|
||||
|
||||
范围:T10/T19/T20/T21/T23 共 5 项修复 + T14 复现证伪,共 6 项。
|
||||
|
||||
| # | 修复内容 | 文件/位置 |
|
||||
|---|---|---|
|
||||
| T10 | 调度器 `await ai_future` 改 `asyncio.wait_for(ai_future, timeout=task_timeout + 30)`——软取消(M3 的 ai.cancel 三层)够不到的卡死路径(代理连接黑洞/滴流输出规避读超时/杀不掉的进程)在 +30s 宽限后硬超时放行,调度循环不再永久停摆;`run_ai` 回调改经 `_set_result/_set_exception_if_pending`(done 检查),wait_for 取消 future 后迟到的 set_result 不再抛 InvalidStateError 打进事件循环 | `aibox.py:318-375` |
|
||||
| T19 | FailoverClient 两处 `except (TransientLLMError, requests.exceptions.RequestException)` 收窄为 `(TransientLLMError,) + _TRANSIENT_NET_ERRORS`(ConnectionError/Timeout/ChunkedEncodingError 显式白名单,模块级常量);非流式 `chat_completion` 内层瞬时上抛清单同步对齐——MissingSchema/InvalidURL 等永久配置错误直接上抛响亮暴露,不再进冷却/转移污染统计与冷却表 | `llm_client.py:36-46,204-206,489,544` |
|
||||
| T20 | `stream_chat` 跟踪 `saw_done`:iter_lines 耗尽但从未收到 `[DONE]`(代理缓冲后提前 FIN/服务端中途断开,末行半截 JSON 被 continue)时抛 `TransientLLMError`——首事件前 FailoverClient 可转移重放,已产出事件则上抛由既有回滚/报错接住,回答无声截断不再被静默视为成功 | `llm_client.py:262-290` |
|
||||
| T21 | 新增 `filter_env_overrides()`(executors 模块级,process_manager 导入):env 覆盖黑名单 `ENV_DENY_EXACT`(PATH/PYTHONPATH/PYTHONHOME/SHELL/IFS/LD_LIBRARY_PATH/PSMODULEPATH/COMSPEC 等)+ `LD_*`/`DYLD_*` 前缀整体拒绝(键大小写不敏感),拒绝键告警留痕;executors 两处与 process_manager 一处 `os.environ.copy().update(env)` 全部接入——LLM(可被提示注入诱导)无法再经 env 参数注入 LD_PRELOAD 等劫持后续进程加载 | `subsystems/executors.py:25-52,329-336,422-429`;`subsystems/process_manager.py:37,163-169` |
|
||||
| T23 | `_validate_message_structure` 双向修复 tool 配对:既有"孤立 tool 消息移除"之外,新增反向——压缩块边界切开后保留侧 assistant 的 tool_calls 失去配对 tool 消息时剥离失配条目(部分失配保留配对项;全部失配删字段,content 为空置空串防非法空消息),OpenAI 协议成对约束不再被压缩破坏导致后续 API 400;顺带删除死代码 `_is_range_safe`(恒 True)与 `_get_protected_indices`(无调用方) | `subsystems/conversation_compressor.py:226-284,334-380` |
|
||||
| T14 | **复现证伪,无需修复**:审计怀疑 FTS5 MATCH 左操作数在 JOIN 别名(`memories_fts f`)场景用错表名导致全程静默降级。实测(SQLite 3.45.1 + trigram):别名场景下 `WHERE memories_fts MATCH ?` 正常命中不报错;端到端 MemoryDatabase 全链路验证——4 字中文词走 FTS trigram 命中、2 字短词走 LIKE 兜底(v3.1 设计行为),`memories_fts` 建表含 `tokenize='trigram'` 确认在用。复现要点:外部内容表(content=memories)的索引需显式 INSERT,且 trigram 需 ≥3 字符 | 复现记录见本节;`subsystems/db.py:196-209`(代码未改动) |
|
||||
|
||||
### 24.8 测试基线(Linux,`aibox/venv`)
|
||||
|
||||
波次 4 后:**584 passed / 1 failed**(31.9s)。`tests/test_n_fixes.py` 追加 18 项(T19 收窄 4 + T20 不完整流 3 + T21 黑名单 4 + T23 配对修复 5 + T10 调度器硬超时 2);1 失败仍为 §11.3 所列 Windows 向用例(`test_execute_timeout`,本机无 ping),零回归。
|
||||
|
||||
### 24.9 波次 5 落地(权限隔离批,2026-08-21)
|
||||
|
||||
范围:T22 + T-L17/T-L18 修复,T1/T2 经产品决策收束(决策:T1 仅加告警不加鉴权、T2 维持现状——本地单用户可信环境定位,接受风险),共 5 项。
|
||||
|
||||
**产品决策记录**(2026-08-21,用户裁决):
|
||||
- **T1**(WS 无鉴权):不加认证,仅绑定非回环地址时启动显著告警(重复 3 次提醒确认网关/WAF 接入)。理由:本地单用户可信环境,引入 token 配置成本大于收益;网络暴露场景由 gateway/WAF 边界承担
|
||||
- **T2**(set_proxy 全局持久化开关):维持现状不改行为,接受该风险。理由:本机可信环境下用户依赖对话中调代理的便利性
|
||||
|
||||
| # | 修复内容 | 文件/位置 |
|
||||
|---|---|---|
|
||||
| T22a | 新增 `TOOL_PROFILE_TASK`:定时任务 AI 显式传精简 profile——剔除 `update_identity`/`update_soul`(改写全局人格提示词)与 `set_proxy`(全局网络面),保留 memory/memo/子代理派生/技能等其余管理面工具;`WORKER_ONLY_TOOLS` 误调拦截条件扩展到 task profile;被污染的任务内容(备忘录标题/正文可能源自网页内容)无法在无人值守时段改写人格与代理 | `aibox.py:1012-1017,104-127,1232-1250,1290-1293,1465` |
|
||||
| T22b | ProcessManager 属主隔离:`start(owner=)` 记录属主,`write`/`kill` 新增 owner 校验(条目与调用方都携带且不一致即拒绝——他人 stdin 不容注入、防提示注入跨会话杀 dev server);任一侧未标识(None)不隔离保持兼容;观察面(list/get_output)不隔离且 `list_procs` 暴露 owner 字段;shutdown/TTL 清理属管理器特权不走校验。`DeepSeekChat` 新增 `owner_id` 构造参数,worker 注册的 `bg_run`/`bg_write`/`bg_kill` 携带(任务 AI 为 `task-<id>`) | `subsystems/process_manager.py:328-345,489-500,597-612,741`;`tools/process_tools.py:81-86,141-142,161-167,222-228,366-371,410`;`aibox.py:1063-1067,1230-1234` |
|
||||
| T1 | 绑定非回环地址(非 localhost/127.0.0.1/::1)时启动告警重复 3 次:任何能建立 TCP 连接的客户端都可使用本智能体(含命令执行能力),提醒确认可信网络/网关接入 | `aibox.py:3170-3178` |
|
||||
| T-L17 | 添加档案时 Key 环境变量名默认值改为按 api_url 域名推荐(deepseek→DEEPSEEK_API_KEY、moonshot→MOONSHOT_API_KEY、aliyuncs→DASHSCOPE_API_KEY 等 9 厂商映射);未知域名(如本地 Ollama)留空且输入为空时取消而非静默变成免密——一路回车不再把 DeepSeek 凭据发给其他厂商 | `model_config.py:285-308,345-352` |
|
||||
| T-L18 | `list_profiles` 展示层输出统一过 `_sanitize_display()`(剥离控制字符)——手写 config.json 嵌 ANSI 转义(`\x1b]0;pwned\x07`)不再于每次列表触发;写入侧 isprintable 校验与读出侧过滤对齐 | `model_config.py:311-328` |
|
||||
|
||||
### 24.10 测试基线(Linux,`aibox/venv`)
|
||||
|
||||
波次 5 后:**598 passed / 1 failed**(38.5s)。`tests/test_n_fixes.py` 追加 14 项(T22 owner 隔离 6 + task profile 4 + T-L17 域名推荐 2 + T-L18 过滤 2);1 失败仍为 §11.3 所列 Windows 向用例(`test_execute_timeout`,本机无 ping),零回归。
|
||||
|
||||
### 24.11 Batch N 收束
|
||||
|
||||
第八轮审计 T 系列 28 项全部处置完毕:波次 1 速赢 7 项 + 波次 2 并发正确性 6 项 + 波次 3 平台兼容 5 项 + 波次 4 LLM 健壮性 5 项 + T14 复现证伪 1 项 + 波次 5 权限隔离 3 项修复与 2 项产品决策收束(T1 告警落地、T2 维持现状)。低危 T-L 系列随批落地 T-L17/T-L18;其余 T-L 项(T-L7~T-L16 等信息级)未在本批处理。
|
||||
@@ -0,0 +1,73 @@
|
||||
# gomsg — pymsg / libmsg 协议的 Go 实现
|
||||
|
||||
Go 版消息通信库,对应 Python [pymsg](../pymsg/) 与 C++ [libmsg](../libmsg/),线协议字节级兼容:
|
||||
|
||||
```
|
||||
[4 字节 big-endian uint32 长度][UTF-8 载荷] (单条上限 64 MiB)
|
||||
```
|
||||
|
||||
## 模块组成
|
||||
|
||||
| 文件 | 内容 | 对应 |
|
||||
|------|------|------|
|
||||
| `frame.go` | 帧编解码:`Encode` / `WriteFrame` / `ReadFrame`(长度上限防护) | pymsg `send`/`recv` |
|
||||
| `client.go` | `MsgClient`:复用长连接 + 请求-响应 + 失败重连一次 | pymsg `MsgClient` |
|
||||
| `server.go` | `MsgServer`:accept 循环 + 每连接 goroutine 串行处理 | pymsg `MsgServer` |
|
||||
| `rpc.go` | `RpcClient` / `RpcServer` / `RpcError`:`_route` 分发、错误脱敏、panic 恢复 | pymsg `rpc.py` |
|
||||
| `asyncqueue.go` | `AsyncQueue[T]`:批量 + 定时 flush,满队列丢新数据(fail-safe) | libmsg `AsyncQueue` |
|
||||
| `logclient.go` | `LogClient`:异步批量推送 db2 Logs 表(级别过滤 + fail-safe) | waf `internal/logclient`、`log_client.py`、liblogc |
|
||||
|
||||
## 使用
|
||||
|
||||
```go
|
||||
// RPC 调用(如调 db2 execute 路由)
|
||||
rpc := gomsg.NewRpcClient("127.0.0.1:8005", 5*time.Second)
|
||||
data, err := rpc.Call("execute", map[string]any{
|
||||
"action": "select", "table": "Users",
|
||||
})
|
||||
// data 为 any:数字是 json.Number(保 int64 精度),需用 .Int64()/.Float64() 转换
|
||||
|
||||
// 提供 RPC 服务
|
||||
srv := gomsg.NewRpcServer("127.0.0.1:8009")
|
||||
srv.AddRoute("health", func(req map[string]any) (any, error) { return map[string]any{"ok": true}, nil })
|
||||
if err := srv.Start(); err != nil { ... }
|
||||
defer srv.Stop()
|
||||
|
||||
// 集中式日志(写入 db2 Logs 表)
|
||||
lc := gomsg.NewLogClient("127.0.0.1:8005")
|
||||
lc.SetLevel("info")
|
||||
lc.Info("backend", "服务启动", nil)
|
||||
defer lc.Flush() // 进程退出前 flush
|
||||
```
|
||||
|
||||
## 与 Python / C++ 的兼容约定
|
||||
|
||||
- **请求**:`{"_route": "<路由名>", ...业务字段}`(`_route` 在分发前被弹出,handler 看不到)
|
||||
- **响应**:`{"success": true, "data": ...}` / `{"success": false, "error": "..."}`
|
||||
- **未知路由**:`{"success": false, "error": "未知路由: <route>"}`
|
||||
- **非法 JSON**:`{"success": false, "error": "invalid JSON request"}`(不回显输入片段)
|
||||
- **错误脱敏**:异常信息中的绝对路径替换为 `<redacted>`(对齐 Python `_sanitize_error`)
|
||||
- **handler panic**:自动恢复并包装为 `success=false`(对应 Python 的 `except Exception`)
|
||||
|
||||
## 安全性
|
||||
|
||||
- 帧长度上限 64 MiB 在读写两侧强制(恶意长度前缀不会导致内存耗尽)
|
||||
- 请求 JSON 解析:Go `encoding/json` 内置 10000 层嵌套上限;数字经 `UseNumber` 以字符串保存,
|
||||
超大数字不会触发 CPU 密集型解析;输入大小受帧上限约束
|
||||
- 异常信息脱敏,不透传绝对路径等内部细节
|
||||
- 数字一律 `json.Number` 保留原始精度,避免 int64 ID 经 float64 丢失
|
||||
|
||||
## 测试
|
||||
|
||||
```bash
|
||||
go test ./... -count=1 -race
|
||||
```
|
||||
|
||||
覆盖:帧往返/超限/截断、Client-Server 回环、自动重连、并发客户端、RPC 成功/失败/脱敏/
|
||||
未知路由/非法 JSON/panic 恢复/超时、AsyncQueue 批量/定时/关停 drain/满队列丢弃、
|
||||
LogClient 对假 db2 服务的推送与级别过滤、db2 不可用时 fail-safe。
|
||||
|
||||
## 来源
|
||||
|
||||
`MsgClient` 连接复用与 `LogClient` 逻辑抽取自 [waf/internal/logclient](../waf/internal/logclient/);
|
||||
API 与协议约定对齐 [pymsg](../pymsg/)(Python)与 [libmsg](../libmsg/)(C++)。
|
||||
@@ -0,0 +1,150 @@
|
||||
package gomsg
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// AsyncQueueConfig 异步批量队列配置(语义对齐 C++ libmsg.AsyncQueueConfig)。
|
||||
type AsyncQueueConfig struct {
|
||||
BatchSize int // 批量发送阈值,达到即 flush;必须 > 0,默认 50
|
||||
FlushInterval time.Duration // 定时 flush 间隔;必须 > 0 且 <= 300s,默认 2s
|
||||
MaxQueueSize int // 队列容量;必须 >= 1,默认 10000;满时丢新数据(对齐 waf logclient 的 fail-safe 语义)
|
||||
}
|
||||
|
||||
func (c AsyncQueueConfig) withDefaults() AsyncQueueConfig {
|
||||
if c.BatchSize <= 0 {
|
||||
c.BatchSize = 50
|
||||
}
|
||||
if c.FlushInterval <= 0 || c.FlushInterval > 300*time.Second {
|
||||
c.FlushInterval = 2 * time.Second
|
||||
}
|
||||
if c.MaxQueueSize < 1 {
|
||||
c.MaxQueueSize = 10000
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
// AsyncQueue 通用异步批量发送队列(对应 C++ libmsg.AsyncQueue)。
|
||||
// Push 非阻塞;队列满时静默丢弃并计数(fail-safe,不影响业务)。
|
||||
type AsyncQueue[T any] struct {
|
||||
cfg AsyncQueueConfig
|
||||
handler func([]T)
|
||||
ch chan T
|
||||
stop chan struct{}
|
||||
wg sync.WaitGroup
|
||||
|
||||
mu sync.Mutex
|
||||
stopped bool
|
||||
dropped uint64
|
||||
flushed uint64
|
||||
}
|
||||
|
||||
// NewAsyncQueue 创建异步批量队列,handler 在单后台 goroutine 中按批调用。
|
||||
func NewAsyncQueue[T any](handler func([]T), cfg AsyncQueueConfig) *AsyncQueue[T] {
|
||||
if handler == nil {
|
||||
panic("gomsg: AsyncQueue handler must not be nil")
|
||||
}
|
||||
cfg = cfg.withDefaults()
|
||||
q := &AsyncQueue[T]{
|
||||
cfg: cfg,
|
||||
handler: handler,
|
||||
ch: make(chan T, cfg.MaxQueueSize),
|
||||
stop: make(chan struct{}),
|
||||
}
|
||||
q.wg.Add(1)
|
||||
go q.loop()
|
||||
return q
|
||||
}
|
||||
|
||||
// Push 入队一条数据;队列满时返回 false(该条被丢弃)。
|
||||
func (q *AsyncQueue[T]) Push(item T) bool {
|
||||
q.mu.Lock()
|
||||
if q.stopped {
|
||||
q.mu.Unlock()
|
||||
return false
|
||||
}
|
||||
q.mu.Unlock()
|
||||
select {
|
||||
case q.ch <- item:
|
||||
return true
|
||||
default:
|
||||
q.mu.Lock()
|
||||
q.dropped++
|
||||
q.mu.Unlock()
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// Shutdown 停止队列,flush 剩余数据后退出后台 goroutine。
|
||||
func (q *AsyncQueue[T]) Shutdown() {
|
||||
q.mu.Lock()
|
||||
if q.stopped {
|
||||
q.mu.Unlock()
|
||||
return
|
||||
}
|
||||
q.stopped = true
|
||||
q.mu.Unlock()
|
||||
close(q.stop)
|
||||
q.wg.Wait()
|
||||
}
|
||||
|
||||
// Stats 返回丢弃与已发送批次数(供监控/测试)。
|
||||
func (q *AsyncQueue[T]) Stats() (dropped, flushed uint64) {
|
||||
q.mu.Lock()
|
||||
defer q.mu.Unlock()
|
||||
return q.dropped, q.flushed
|
||||
}
|
||||
|
||||
func (q *AsyncQueue[T]) loop() {
|
||||
defer q.wg.Done()
|
||||
batch := make([]T, 0, q.cfg.BatchSize)
|
||||
ticker := time.NewTicker(q.cfg.FlushInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case item := <-q.ch:
|
||||
batch = append(batch, item)
|
||||
// 非阻塞续取,直到凑满 batchSize 或暂时无数据
|
||||
for len(batch) < q.cfg.BatchSize {
|
||||
select {
|
||||
case item := <-q.ch:
|
||||
batch = append(batch, item)
|
||||
default:
|
||||
goto send
|
||||
}
|
||||
}
|
||||
send:
|
||||
if len(batch) >= q.cfg.BatchSize {
|
||||
q.send(batch)
|
||||
batch = batch[:0]
|
||||
}
|
||||
case <-ticker.C:
|
||||
if len(batch) > 0 {
|
||||
q.send(batch)
|
||||
batch = batch[:0]
|
||||
}
|
||||
case <-q.stop:
|
||||
// 退出前 drain 剩余数据
|
||||
for {
|
||||
select {
|
||||
case item := <-q.ch:
|
||||
batch = append(batch, item)
|
||||
default:
|
||||
if len(batch) > 0 {
|
||||
q.send(batch)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (q *AsyncQueue[T]) send(batch []T) {
|
||||
q.handler(batch)
|
||||
q.mu.Lock()
|
||||
q.flushed++
|
||||
q.mu.Unlock()
|
||||
}
|
||||
@@ -0,0 +1,120 @@
|
||||
package gomsg
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func collectBatches(batches *[][]int, mu *sync.Mutex) func([]int) {
|
||||
return func(b []int) {
|
||||
mu.Lock()
|
||||
*batches = append(*batches, append([]int(nil), b...))
|
||||
mu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
func TestAsyncQueueFlushOnBatchSize(t *testing.T) {
|
||||
var batches [][]int
|
||||
var mu sync.Mutex
|
||||
q := NewAsyncQueue[int](collectBatches(&batches, &mu), AsyncQueueConfig{
|
||||
BatchSize: 5,
|
||||
FlushInterval: time.Hour, // 只测批量阈值,不测定时
|
||||
})
|
||||
for i := 1; i <= 5; i++ {
|
||||
if !q.Push(i) {
|
||||
t.Fatal("push should succeed")
|
||||
}
|
||||
}
|
||||
deadline := time.Now().Add(2 * time.Second)
|
||||
for {
|
||||
mu.Lock()
|
||||
n := len(batches)
|
||||
mu.Unlock()
|
||||
if n >= 1 {
|
||||
break
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatal("batch not flushed in time")
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if len(batches[0]) != 5 || batches[0][4] != 5 {
|
||||
t.Fatalf("batch = %v", batches[0])
|
||||
}
|
||||
q.Shutdown()
|
||||
}
|
||||
|
||||
func TestAsyncQueueFlushOnInterval(t *testing.T) {
|
||||
var batches [][]int
|
||||
var mu sync.Mutex
|
||||
q := NewAsyncQueue[int](collectBatches(&batches, &mu), AsyncQueueConfig{
|
||||
BatchSize: 100,
|
||||
FlushInterval: 100 * time.Millisecond,
|
||||
})
|
||||
if !q.Push(42) {
|
||||
t.Fatal("push should succeed")
|
||||
}
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
mu.Lock()
|
||||
n := len(batches)
|
||||
mu.Unlock()
|
||||
if n < 1 {
|
||||
t.Fatal("expected interval flush")
|
||||
}
|
||||
q.Shutdown()
|
||||
}
|
||||
|
||||
func TestAsyncQueueShutdownDrains(t *testing.T) {
|
||||
var batches [][]int
|
||||
var mu sync.Mutex
|
||||
q := NewAsyncQueue[int](collectBatches(&batches, &mu), AsyncQueueConfig{
|
||||
BatchSize: 100,
|
||||
FlushInterval: time.Hour,
|
||||
})
|
||||
for i := 1; i <= 3; i++ {
|
||||
q.Push(i)
|
||||
}
|
||||
q.Shutdown()
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if len(batches) != 1 || len(batches[0]) != 3 {
|
||||
t.Fatalf("expected drained batch [1 2 3], got %v", batches)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAsyncQueueDropOnFull(t *testing.T) {
|
||||
var batches [][]int
|
||||
var mu sync.Mutex
|
||||
q := NewAsyncQueue[int](collectBatches(&batches, &mu), AsyncQueueConfig{
|
||||
BatchSize: 1000,
|
||||
FlushInterval: time.Hour,
|
||||
MaxQueueSize: 4,
|
||||
})
|
||||
// 后台协程可能已取走部分,填满队列直到 Push 返回 false
|
||||
dropped := false
|
||||
for i := 0; i < 100; i++ {
|
||||
if !q.Push(i) {
|
||||
dropped = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !dropped {
|
||||
t.Fatal("expected drop when queue full")
|
||||
}
|
||||
d, _ := q.Stats()
|
||||
if d == 0 {
|
||||
t.Fatal("dropped counter should be > 0")
|
||||
}
|
||||
q.Shutdown()
|
||||
}
|
||||
|
||||
func TestAsyncQueueRejectAfterShutdown(t *testing.T) {
|
||||
q := NewAsyncQueue[int](func([]int) {}, AsyncQueueConfig{})
|
||||
q.Shutdown()
|
||||
if q.Push(1) {
|
||||
t.Fatal("push after shutdown should fail")
|
||||
}
|
||||
}
|
||||
+106
@@ -0,0 +1,106 @@
|
||||
package gomsg
|
||||
|
||||
import (
|
||||
"net"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// DefaultTimeout 单次请求-响应交互的总超时。
|
||||
const DefaultTimeout = 5 * time.Second
|
||||
|
||||
// MsgClient pymsg 协议客户端:复用长连接 + 请求-响应调用。
|
||||
// 连接断开时自动重建一次(对应 Python pymsg.MsgClient / C++ libmsg.MsgClient)。
|
||||
type MsgClient struct {
|
||||
addr string
|
||||
timeout time.Duration
|
||||
|
||||
mu sync.Mutex
|
||||
conn net.Conn
|
||||
}
|
||||
|
||||
// NewMsgClient 创建客户端。timeout 为 0 时使用 DefaultTimeout。
|
||||
func NewMsgClient(addr string, timeout time.Duration) *MsgClient {
|
||||
if timeout <= 0 {
|
||||
timeout = DefaultTimeout
|
||||
}
|
||||
return &MsgClient{addr: addr, timeout: timeout}
|
||||
}
|
||||
|
||||
// Call 发送一条请求并读取响应(复用长连接,失败时关闭重连一次)。
|
||||
// 返回响应载荷;超时 / 连接错误返回 error。
|
||||
func (c *MsgClient) Call(payload []byte) ([]byte, error) {
|
||||
resp, err := c.call(payload)
|
||||
if err == nil {
|
||||
return resp, nil
|
||||
}
|
||||
// 旧连接可能已断开(服务端重启等),关闭后重连一次
|
||||
c.closeConn()
|
||||
return c.call(payload)
|
||||
}
|
||||
|
||||
// Send 发送一条请求,不等待响应(fire-and-forget)。失败时重连一次。
|
||||
func (c *MsgClient) Send(payload []byte) error {
|
||||
if err := c.sendOnly(payload); err != nil {
|
||||
c.closeConn()
|
||||
return c.sendOnly(payload)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Close 关闭底层连接。
|
||||
func (c *MsgClient) Close() error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return c.closeLocked()
|
||||
}
|
||||
|
||||
func (c *MsgClient) call(payload []byte) ([]byte, error) {
|
||||
conn, err := c.getConn()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
_ = conn.SetDeadline(time.Now().Add(c.timeout))
|
||||
if err := WriteFrame(conn, payload); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return ReadFrame(conn, 0)
|
||||
}
|
||||
|
||||
func (c *MsgClient) sendOnly(payload []byte) error {
|
||||
conn, err := c.getConn()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_ = conn.SetWriteDeadline(time.Now().Add(c.timeout))
|
||||
return WriteFrame(conn, payload)
|
||||
}
|
||||
|
||||
func (c *MsgClient) getConn() (net.Conn, error) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if c.conn != nil {
|
||||
return c.conn, nil
|
||||
}
|
||||
conn, err := net.DialTimeout("tcp", c.addr, c.timeout)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.conn = conn
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
func (c *MsgClient) closeConn() {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
_ = c.closeLocked()
|
||||
}
|
||||
|
||||
func (c *MsgClient) closeLocked() error {
|
||||
if c.conn != nil {
|
||||
err := c.conn.Close()
|
||||
c.conn = nil
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,170 @@
|
||||
package gomsg
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// startEchoServer 启动一个 MsgServer,把收到的载荷原样回复。
|
||||
func startEchoServer(t *testing.T) (*MsgServer, string) {
|
||||
t.Helper()
|
||||
s := NewMsgServer("127.0.0.1:0", func(payload []byte) []byte {
|
||||
return append([]byte("echo:"), payload...)
|
||||
})
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(s.Stop)
|
||||
return s, s.Addr().String()
|
||||
}
|
||||
|
||||
func TestMsgClientCall(t *testing.T) {
|
||||
_, addr := startEchoServer(t)
|
||||
c := NewMsgClient(addr, 0)
|
||||
defer c.Close()
|
||||
|
||||
resp, err := c.Call([]byte(`{"hello":"世界"}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := "echo:{\"hello\":\"世界\"}"
|
||||
if string(resp) != want {
|
||||
t.Fatalf("resp = %q, want %q", resp, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMsgClientReusesConn(t *testing.T) {
|
||||
_, addr := startEchoServer(t)
|
||||
c := NewMsgClient(addr, 0)
|
||||
defer c.Close()
|
||||
|
||||
for i := 0; i < 5; i++ {
|
||||
if _, err := c.Call([]byte("x")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if c.conn == nil {
|
||||
t.Fatal("expected connection reuse")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMsgClientReconnectAfterServerRestart(t *testing.T) {
|
||||
s := NewMsgServer("127.0.0.1:0", func(p []byte) []byte { return p })
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
addr := s.Addr().String()
|
||||
|
||||
c := NewMsgClient(addr, 2*time.Second)
|
||||
defer c.Close()
|
||||
if _, err := c.Call([]byte("before")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// 服务端重启:旧连接失效,Call 应自动重连成功
|
||||
s.Stop()
|
||||
// 重新监听同一地址:旧 listener 已释放端口
|
||||
s3 := NewMsgServer(addr, func(p []byte) []byte { return p })
|
||||
if err := s3.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s3.Stop()
|
||||
|
||||
resp, err := c.Call([]byte("after"))
|
||||
if err != nil {
|
||||
t.Fatalf("expected auto-reconnect, got %v", err)
|
||||
}
|
||||
if string(resp) != "after" {
|
||||
t.Fatalf("resp = %q", resp)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMsgClientSendNoReply(t *testing.T) {
|
||||
got := make(chan []byte, 1)
|
||||
s := NewMsgServer("127.0.0.1:0", func(payload []byte) []byte {
|
||||
got <- payload
|
||||
return nil // 不回复
|
||||
})
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Stop()
|
||||
|
||||
c := NewMsgClient(s.Addr().String(), 0)
|
||||
defer c.Close()
|
||||
if err := c.Send([]byte("fire-and-forget")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
select {
|
||||
case p := <-got:
|
||||
if string(p) != "fire-and-forget" {
|
||||
t.Fatalf("payload = %q", p)
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("server did not receive message")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMsgClientTimeout(t *testing.T) {
|
||||
s := NewMsgServer("127.0.0.1:0", func(payload []byte) []byte {
|
||||
time.Sleep(500 * time.Millisecond) // 拖过客户端超时
|
||||
return nil
|
||||
})
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Stop()
|
||||
|
||||
c := NewMsgClient(s.Addr().String(), 50*time.Millisecond)
|
||||
defer c.Close()
|
||||
if _, err := c.Call([]byte("slow")); err == nil {
|
||||
t.Fatal("expected timeout error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMsgServerStopClosesConns(t *testing.T) {
|
||||
_, addr := startEchoServer(t)
|
||||
c := NewMsgClient(addr, 0)
|
||||
defer c.Close()
|
||||
if _, err := c.Call([]byte("ping")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMsgServerConcurrentClients(t *testing.T) {
|
||||
s := NewMsgServer("127.0.0.1:0", func(payload []byte) []byte {
|
||||
return bytes.ToUpper(payload)
|
||||
})
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Stop()
|
||||
|
||||
errCh := make(chan error, 8)
|
||||
for i := 0; i < 8; i++ {
|
||||
go func(n int) {
|
||||
c := NewMsgClient(s.Addr().String(), 0)
|
||||
defer c.Close()
|
||||
for j := 0; j < 10; j++ {
|
||||
in := []byte{byte('a' + n)}
|
||||
resp, err := c.Call(in)
|
||||
if err != nil {
|
||||
errCh <- err
|
||||
return
|
||||
}
|
||||
if len(resp) != 1 || resp[0] != in[0]-'a'+'A' {
|
||||
errCh <- errors.New("unexpected reply")
|
||||
return
|
||||
}
|
||||
}
|
||||
errCh <- nil
|
||||
}(i)
|
||||
}
|
||||
for i := 0; i < 8; i++ {
|
||||
if err := <-errCh; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
// Package gomsg 是 pymsg(Python)/ libmsg(C++)协议的 Go 实现。
|
||||
//
|
||||
// 线协议(与 Python / C++ 字节级兼容):
|
||||
//
|
||||
// [4 字节 big-endian uint32 长度][UTF-8 载荷]
|
||||
//
|
||||
// 消息大小上限 64 MiB(与 Python pymsg / C++ libmsg 对齐)。
|
||||
//
|
||||
// 覆盖功能:帧编解码(本文件)、MsgClient/MsgServer(client.go / server.go)、
|
||||
// JSON RPC 约定层(rpc.go)、通用异步批量队列(asyncqueue.go)与
|
||||
// 集中式日志客户端(logclient.go,对应各服务目录下的 log_client.py / liblogc)。
|
||||
package gomsg
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"io"
|
||||
)
|
||||
|
||||
// MaxMsgSize 单条消息载荷大小上限,与 Python pymsg / C++ libmsg 保持一致。
|
||||
const MaxMsgSize = 64 * 1024 * 1024
|
||||
|
||||
var headerSize = 4
|
||||
|
||||
// Encode 将载荷封装为 pymsg 帧:4 字节大端长度前缀 + 载荷。
|
||||
func Encode(payload []byte) ([]byte, error) {
|
||||
if len(payload) > MaxMsgSize {
|
||||
return nil, fmt.Errorf("payload too large: %d bytes (max %d)", len(payload), MaxMsgSize)
|
||||
}
|
||||
frame := make([]byte, headerSize+len(payload))
|
||||
binary.BigEndian.PutUint32(frame, uint32(len(payload)))
|
||||
copy(frame[headerSize:], payload)
|
||||
return frame, nil
|
||||
}
|
||||
|
||||
// WriteFrame 将载荷以 pymsg 帧形式写入 w。
|
||||
func WriteFrame(w io.Writer, payload []byte) error {
|
||||
frame, err := Encode(payload)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = w.Write(frame)
|
||||
return err
|
||||
}
|
||||
|
||||
// ReadFrame 从 r 读取一条 pymsg 帧,返回载荷。
|
||||
// max 为单条消息大小上限(传 0 表示使用 MaxMsgSize)。
|
||||
// 长度前缀超过上限视为协议错误(防止恶意长度导致内存耗尽)。
|
||||
func ReadFrame(r io.Reader, max int) ([]byte, error) {
|
||||
if max <= 0 || max > MaxMsgSize {
|
||||
max = MaxMsgSize
|
||||
}
|
||||
header := make([]byte, headerSize)
|
||||
if _, err := io.ReadFull(r, header); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
length := binary.BigEndian.Uint32(header)
|
||||
if length > uint32(max) {
|
||||
return nil, fmt.Errorf("message too large: %d bytes (max %d)", length, max)
|
||||
}
|
||||
payload := make([]byte, length)
|
||||
if _, err := io.ReadFull(r, payload); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return payload, nil
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
package gomsg
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"io"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestEncodeRoundTrip(t *testing.T) {
|
||||
payload := []byte(`{"hello":"世界"}`)
|
||||
frame, err := Encode(payload)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(frame) != 4+len(payload) {
|
||||
t.Fatalf("frame length = %d, want %d", len(frame), 4+len(payload))
|
||||
}
|
||||
if got := binary.BigEndian.Uint32(frame[:4]); got != uint32(len(payload)) {
|
||||
t.Fatalf("header = %d, want %d", got, len(payload))
|
||||
}
|
||||
got, err := ReadFrame(bytes.NewReader(frame), 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !bytes.Equal(got, payload) {
|
||||
t.Fatalf("payload mismatch: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEncodeTooLarge(t *testing.T) {
|
||||
payload := make([]byte, MaxMsgSize+1)
|
||||
if _, err := Encode(payload); err == nil {
|
||||
t.Fatal("expected error for oversized payload")
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadFrameTooLarge(t *testing.T) {
|
||||
header := make([]byte, 4)
|
||||
binary.BigEndian.PutUint32(header, MaxMsgSize+1)
|
||||
_, err := ReadFrame(bytes.NewReader(header), 0)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for oversized header")
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadFrameTruncated(t *testing.T) {
|
||||
frame, _ := Encode([]byte("abcdef"))
|
||||
if _, err := ReadFrame(bytes.NewReader(frame[:len(frame)-3]), 0); err == nil {
|
||||
t.Fatal("expected error for truncated payload")
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadFrameMaxParam(t *testing.T) {
|
||||
frame, _ := Encode([]byte("hello"))
|
||||
if _, err := ReadFrame(bytes.NewReader(frame), 4); err == nil {
|
||||
t.Fatal("expected error when max < payload length")
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadFrameEOF(t *testing.T) {
|
||||
if _, err := ReadFrame(bytes.NewReader(nil), 0); !errors.Is(err, io.EOF) {
|
||||
t.Fatalf("expected EOF, got %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
module github.com/micro_servers/gomsg
|
||||
|
||||
go 1.22
|
||||
@@ -0,0 +1,173 @@
|
||||
package gomsg
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"time"
|
||||
)
|
||||
|
||||
// 集中式日志客户端:通过 pymsg/TCP 调用 db2(或 db)execute 路由,
|
||||
// 逐条 INSERT 到 Logs 表。对应各服务目录下的 log_client.py 与 C++ liblogc。
|
||||
//
|
||||
// 设计要点(与 waf internal/logclient / Python log_client.py 对齐):
|
||||
// - 异步非阻塞:AsyncQueue 批量推送,不拖慢业务
|
||||
// - fail-safe:db 服务不可用时静默丢弃,绝不影响业务
|
||||
// - 缓冲批量:满 50 条或超时 2 秒自动 flush
|
||||
|
||||
const defaultDBAddr = "127.0.0.1:8005"
|
||||
|
||||
// 日志级别数值(越小越详细)。
|
||||
const (
|
||||
levelDebug = 0
|
||||
levelInfo = 1
|
||||
levelWarn = 2
|
||||
levelError = 3
|
||||
)
|
||||
|
||||
var levelMap = map[string]int{
|
||||
"debug": levelDebug, "DEBUG": levelDebug,
|
||||
"info": levelInfo, "INFO": levelInfo,
|
||||
"warn": levelWarn, "warning": levelWarn, "WARN": levelWarn, "WARNING": levelWarn,
|
||||
"error": levelError, "ERROR": levelError,
|
||||
}
|
||||
|
||||
// Entry 单条日志。
|
||||
type Entry struct {
|
||||
Service string
|
||||
Level string
|
||||
Message string
|
||||
TraceID string
|
||||
Extra map[string]any
|
||||
}
|
||||
|
||||
// LogClient 异步批量日志推送客户端。
|
||||
type LogClient struct {
|
||||
queue *AsyncQueue[Entry]
|
||||
client *MsgClient
|
||||
minLevel int
|
||||
}
|
||||
|
||||
// NewLogClient 创建日志客户端并启动后台批量推送。
|
||||
// dbAddr 为空时依次取环境变量 LOG_DB_ADDR、默认 127.0.0.1:8005;
|
||||
// 默认最低级别 info,可用 SetLevel 调整(环境变量 LOG_LEVEL 同样生效)。
|
||||
func NewLogClient(dbAddr string) *LogClient {
|
||||
if dbAddr == "" {
|
||||
dbAddr = os.Getenv("LOG_DB_ADDR")
|
||||
}
|
||||
if dbAddr == "" {
|
||||
dbAddr = defaultDBAddr
|
||||
}
|
||||
lc := &LogClient{
|
||||
client: NewMsgClient(dbAddr, 5*time.Second),
|
||||
minLevel: levelInfo,
|
||||
}
|
||||
lc.queue = NewAsyncQueue[Entry](lc.sendBatch, AsyncQueueConfig{})
|
||||
if lv := os.Getenv("LOG_LEVEL"); lv != "" {
|
||||
lc.SetLevel(lv)
|
||||
}
|
||||
return lc
|
||||
}
|
||||
|
||||
// SetLevel 设置最低日志级别(debug|info|warn|error,大小写不敏感)。
|
||||
func (c *LogClient) SetLevel(level string) {
|
||||
if n, ok := levelMap[level]; ok {
|
||||
c.minLevel = n
|
||||
}
|
||||
}
|
||||
|
||||
// Push 推送一条日志(非阻塞,队列满或低于最低级别时静默丢弃)。
|
||||
func (c *LogClient) Push(e Entry) bool {
|
||||
n, ok := levelMap[e.Level]
|
||||
if !ok || n < c.minLevel {
|
||||
return false
|
||||
}
|
||||
return c.queue.Push(e)
|
||||
}
|
||||
|
||||
// Debug 推送 DEBUG 级别日志。
|
||||
func (c *LogClient) Debug(service, message string, extra map[string]any) bool {
|
||||
return c.Push(Entry{Service: service, Level: "DEBUG", Message: message, Extra: extra})
|
||||
}
|
||||
|
||||
// Info 推送 INFO 级别日志。
|
||||
func (c *LogClient) Info(service, message string, extra map[string]any) bool {
|
||||
return c.Push(Entry{Service: service, Level: "INFO", Message: message, Extra: extra})
|
||||
}
|
||||
|
||||
// Warning 推送 WARNING 级别日志。
|
||||
func (c *LogClient) Warning(service, message string, extra map[string]any) bool {
|
||||
return c.Push(Entry{Service: service, Level: "WARNING", Message: message, Extra: extra})
|
||||
}
|
||||
|
||||
// Error 推送 ERROR 级别日志。
|
||||
func (c *LogClient) Error(service, message string, extra map[string]any) bool {
|
||||
return c.Push(Entry{Service: service, Level: "ERROR", Message: message, Extra: extra})
|
||||
}
|
||||
|
||||
// Flush 停止队列并 flush 剩余日志(用于进程退出前)。
|
||||
func (c *LogClient) Flush() {
|
||||
c.queue.Shutdown()
|
||||
}
|
||||
|
||||
// dbV2Request execute 路由的 INSERT 请求体(与 Python/C++ 各 log 客户端一致)。
|
||||
type dbV2Request struct {
|
||||
Route string `json:"_route"`
|
||||
Action string `json:"action"`
|
||||
Table string `json:"table"`
|
||||
Values map[string]any `json:"values"`
|
||||
}
|
||||
|
||||
// dbV2Response execute 路由的通用响应。
|
||||
type dbV2Response struct {
|
||||
Success bool `json:"success"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
func (c *LogClient) sendBatch(batch []Entry) {
|
||||
now := time.Now().UTC().Format(time.RFC3339)
|
||||
for _, e := range batch {
|
||||
c.sendOne(e, now)
|
||||
}
|
||||
}
|
||||
|
||||
// sendOne 通过 pymsg TCP 协议发送单条日志到 db2(fail-safe:任何错误静默丢弃)。
|
||||
func (c *LogClient) sendOne(e Entry, ts string) {
|
||||
extraJSON := ""
|
||||
if e.Extra != nil {
|
||||
b, err := json.Marshal(e.Extra)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
extraJSON = string(b)
|
||||
}
|
||||
|
||||
values := map[string]any{
|
||||
"ts": ts,
|
||||
"service": e.Service,
|
||||
"level": e.Level,
|
||||
"message": e.Message,
|
||||
}
|
||||
if e.TraceID != "" {
|
||||
values["trace_id"] = e.TraceID
|
||||
}
|
||||
if extraJSON != "" {
|
||||
values["extra"] = extraJSON
|
||||
}
|
||||
|
||||
req := dbV2Request{Route: "execute", Action: "insert", Table: "Logs", Values: values}
|
||||
data, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
resp, err := c.client.Call(data)
|
||||
if err != nil {
|
||||
return // fail-safe
|
||||
}
|
||||
|
||||
var r dbV2Response
|
||||
if json.Unmarshal(resp, &r) != nil || !r.Success {
|
||||
fmt.Fprintf(os.Stderr, "[gomsg:logclient] INSERT failed: %s\n", r.Error)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
package gomsg
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// fakeDB2 模拟 db2 的 execute 路由:校验 Logs INSERT 请求格式并返回成功。
|
||||
type fakeDB2 struct {
|
||||
server *MsgServer
|
||||
mu sync.Mutex
|
||||
inserts []map[string]any
|
||||
failures int
|
||||
}
|
||||
|
||||
func newFakeDB2(t *testing.T) *fakeDB2 {
|
||||
f := &fakeDB2{}
|
||||
f.server = NewMsgServer("127.0.0.1:0", func(payload []byte) []byte {
|
||||
var req map[string]any
|
||||
if err := json.Unmarshal(payload, &req); err != nil {
|
||||
return []byte(`{"success":false,"error":"invalid JSON request"}`)
|
||||
}
|
||||
// db2 真实语义:非法 JSON 返回 success=false;其余按通用响应
|
||||
if req["_route"] != "execute" {
|
||||
return []byte(`{"success":false,"error":"未知路由"}`)
|
||||
}
|
||||
f.mu.Lock()
|
||||
f.inserts = append(f.inserts, req)
|
||||
f.mu.Unlock()
|
||||
return []byte(`{"success":true,"data":{}}`)
|
||||
})
|
||||
if err := f.server.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(f.server.Stop)
|
||||
return f
|
||||
}
|
||||
|
||||
func (f *fakeDB2) count() int {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
return len(f.inserts)
|
||||
}
|
||||
|
||||
func TestLogClientPushesToDB2(t *testing.T) {
|
||||
db := newFakeDB2(t)
|
||||
lc := NewLogClient(db.server.Addr().String())
|
||||
defer lc.Flush()
|
||||
|
||||
lc.SetLevel("debug")
|
||||
lc.Debug("test-svc", "debug msg", nil)
|
||||
lc.Info("test-svc", "hello", map[string]any{"k": "v"})
|
||||
lc.Error("test-svc", "boom", nil)
|
||||
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for db.count() < 3 {
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatal("logs not pushed in time")
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
|
||||
db.mu.Lock()
|
||||
defer db.mu.Unlock()
|
||||
first := db.inserts[0]
|
||||
if first["_route"] != "execute" {
|
||||
t.Fatalf("_route = %v", first["_route"])
|
||||
}
|
||||
vals := first["values"].(map[string]any)
|
||||
if vals["service"] != "test-svc" || vals["level"] != "DEBUG" || vals["message"] != "debug msg" {
|
||||
t.Fatalf("values = %v", vals)
|
||||
}
|
||||
// ts 必须是 ISO 8601 字符串
|
||||
if _, ok := vals["ts"].(string); !ok {
|
||||
t.Fatalf("ts missing: %v", vals)
|
||||
}
|
||||
// 第二条(Info)的 extra 应序列化为 JSON 字符串
|
||||
vals2 := db.inserts[1]["values"].(map[string]any)
|
||||
if vals2["extra"] != `{"k":"v"}` {
|
||||
t.Fatalf("extra = %v", vals2["extra"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestLogClientLevelFilter(t *testing.T) {
|
||||
db := newFakeDB2(t)
|
||||
lc := NewLogClient(db.server.Addr().String())
|
||||
defer lc.Flush()
|
||||
|
||||
lc.Info("svc", "visible", nil)
|
||||
lc.Debug("svc", "hidden", nil) // 默认 minLevel=info,debug 应被过滤
|
||||
lc.Push(Entry{Service: "svc", Level: "NOPE", Message: "hidden"})
|
||||
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for db.count() < 1 {
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatal("expected exactly 1 log")
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
time.Sleep(100 * time.Millisecond) // 确认没有多余的推送
|
||||
if db.count() != 1 {
|
||||
t.Fatalf("expected 1 insert, got %d", db.count())
|
||||
}
|
||||
}
|
||||
|
||||
func TestLogClientFailSafeWhenDBDown(t *testing.T) {
|
||||
lc := NewLogClient("127.0.0.1:1") // 不可达地址
|
||||
defer lc.Flush()
|
||||
lc.Info("svc", "will be dropped", nil)
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
// 不 panic、不阻塞即视为通过(fail-safe 语义)
|
||||
}
|
||||
+224
@@ -0,0 +1,224 @@
|
||||
package gomsg
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// JSON RPC 约定层(与 Python pymsg.rpc / C++ libmsg RpcClient/RpcServer 兼容):
|
||||
//
|
||||
// 请求: {"_route": "<路由名>", ...业务字段}
|
||||
// 响应: {"success": true, "data": ...}
|
||||
// {"success": false, "error": "..."}
|
||||
//
|
||||
// 数字一律以 json.Number 保留原始精度(避免 int64 经 float64 丢失)。
|
||||
|
||||
// RpcError 远端返回 success=false 的业务错误,或响应不符合 RPC 约定。
|
||||
type RpcError struct {
|
||||
msg string
|
||||
}
|
||||
|
||||
func (e *RpcError) Error() string { return e.msg }
|
||||
|
||||
// RpcHandler 处理一条 RPC 请求,返回值放入 data,错误包装为 success=false。
|
||||
type RpcHandler func(req map[string]any) (any, error)
|
||||
|
||||
// RpcClient JSON RPC 客户端(复用 MsgClient 长连接)。
|
||||
type RpcClient struct {
|
||||
client *MsgClient
|
||||
}
|
||||
|
||||
// NewRpcClient 创建客户端。timeout 为 0 时使用 DefaultTimeout。
|
||||
func NewRpcClient(addr string, timeout time.Duration) *RpcClient {
|
||||
return &RpcClient{client: NewMsgClient(addr, timeout)}
|
||||
}
|
||||
|
||||
// Call 调用指定路由,返回响应中的 data 字段(可能为 nil)。
|
||||
// success=false 或非法响应时返回 *RpcError。
|
||||
func (c *RpcClient) Call(route string, payload map[string]any) (any, error) {
|
||||
req := make(map[string]any, len(payload)+1)
|
||||
for k, v := range payload {
|
||||
req[k] = v
|
||||
}
|
||||
req["_route"] = route
|
||||
body, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
resp, err := c.client.Call(body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var obj map[string]any
|
||||
dec := json.NewDecoder(bytes.NewReader(resp))
|
||||
dec.UseNumber()
|
||||
if err := dec.Decode(&obj); err != nil || obj == nil {
|
||||
return nil, &RpcError{msg: fmt.Sprintf("响应不是合法 JSON: %v", err)}
|
||||
}
|
||||
success, ok := obj["success"].(bool)
|
||||
if !ok {
|
||||
return nil, &RpcError{msg: fmt.Sprintf("非法 RPC 响应: %s", truncate(resp))}
|
||||
}
|
||||
if !success {
|
||||
errMsg, _ := obj["error"].(string)
|
||||
return nil, &RpcError{msg: errMsg}
|
||||
}
|
||||
return obj["data"], nil
|
||||
}
|
||||
|
||||
// Close 关闭底层连接。
|
||||
func (c *RpcClient) Close() error { return c.client.Close() }
|
||||
|
||||
// RpcServer JSON RPC 服务端:基于 MsgServer 的路由分发与响应包装。
|
||||
type RpcServer struct {
|
||||
*MsgServer
|
||||
routes map[string]RpcHandler
|
||||
mu sync.RWMutex
|
||||
// OnError 可选钩子:handler 抛错时调用(用于记日志)。
|
||||
OnError func(route string, err error)
|
||||
// JSONMarshal 响应序列化函数,默认 json.Marshal(对应 Python 的 json_default)。
|
||||
JSONMarshal func(v any) ([]byte, error)
|
||||
}
|
||||
|
||||
// NewRpcServer 创建 RPC 服务端(不监听,需调用 Start)。
|
||||
func NewRpcServer(addr string) *RpcServer {
|
||||
s := &RpcServer{routes: make(map[string]RpcHandler)}
|
||||
s.MsgServer = NewMsgServer(addr, s.dispatch)
|
||||
return s
|
||||
}
|
||||
|
||||
// AddRoute 注册路由处理函数。
|
||||
func (s *RpcServer) AddRoute(name string, fn RpcHandler) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.routes[name] = fn
|
||||
}
|
||||
|
||||
// dispatch 解析请求、查找路由、包装响应;任何输入错误都不向客户端回显内部细节。
|
||||
// handler 内 panic 会被恢复并包装为 success=false(对应 Python 的 except Exception)。
|
||||
func (s *RpcServer) dispatch(payload []byte) (reply []byte) {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
err := fmt.Errorf("handler panic: %v", r)
|
||||
if s.OnError != nil {
|
||||
s.OnError("", err)
|
||||
}
|
||||
reply = s.reply(false, nil, sanitizeError(err.Error()))
|
||||
}
|
||||
}()
|
||||
|
||||
req, err := decodeRequest(payload)
|
||||
if err != nil {
|
||||
return s.reply(false, nil, "invalid JSON request")
|
||||
}
|
||||
|
||||
routeVal, has := req["_route"]
|
||||
delete(req, "_route")
|
||||
route := ""
|
||||
switch r := routeVal.(type) {
|
||||
case string:
|
||||
route = r
|
||||
case nil:
|
||||
// 保持 Python 行为:缺省 _route 时按空路由查找(未知路由)
|
||||
default:
|
||||
if has {
|
||||
route = fmt.Sprintf("%v", r)
|
||||
}
|
||||
}
|
||||
|
||||
s.mu.RLock()
|
||||
handler := s.routes[route]
|
||||
s.mu.RUnlock()
|
||||
if handler == nil {
|
||||
return s.reply(false, nil, "未知路由: "+route)
|
||||
}
|
||||
|
||||
data, err := handler(req)
|
||||
if err != nil {
|
||||
if s.OnError != nil {
|
||||
s.OnError(route, err)
|
||||
}
|
||||
return s.reply(false, nil, sanitizeError(err.Error()))
|
||||
}
|
||||
return s.reply(true, data, "")
|
||||
}
|
||||
|
||||
func (s *RpcServer) reply(success bool, data any, errMsg string) []byte {
|
||||
resp := map[string]any{"success": success}
|
||||
if data != nil {
|
||||
resp["data"] = data
|
||||
}
|
||||
if errMsg != "" {
|
||||
resp["error"] = errMsg
|
||||
}
|
||||
marshal := s.JSONMarshal
|
||||
if marshal == nil {
|
||||
marshal = json.Marshal
|
||||
}
|
||||
out, err := marshal(resp)
|
||||
if err != nil {
|
||||
out, _ = json.Marshal(map[string]any{
|
||||
"success": false,
|
||||
"error": sanitizeError(err.Error()),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// decodeRequest 解析请求 JSON:UseNumber 保留整数精度;
|
||||
// Go 的 encoding/json 内置 10000 层嵌套上限,数字位数字符串化保存,
|
||||
// 输入大小又受帧协议 64 MiB 上限约束,天然规避深度/超大数字 DoS。
|
||||
func decodeRequest(payload []byte) (map[string]any, error) {
|
||||
dec := json.NewDecoder(bytes.NewReader(payload))
|
||||
dec.UseNumber()
|
||||
var req map[string]any
|
||||
if err := dec.Decode(&req); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if req == nil {
|
||||
return nil, fmt.Errorf("empty request")
|
||||
}
|
||||
return req, nil
|
||||
}
|
||||
|
||||
// pathRe 匹配绝对路径(如 /opt/micro_servers/db/xxx.py),用于脱敏异常信息。
|
||||
var pathRe = regexp.MustCompile(`/[A-Za-z0-9_.\-]+(?:/[A-Za-z0-9_.\-]+)*`)
|
||||
|
||||
// sanitizeError 脱敏异常信息中的绝对路径等内部细节,保留业务错误语义
|
||||
// (对齐 Python pymsg._sanitize_error:前导字符为单词字符时不匹配)。
|
||||
func sanitizeError(msg string) string {
|
||||
if msg == "" {
|
||||
return msg
|
||||
}
|
||||
var b strings.Builder
|
||||
last := 0
|
||||
for _, idx := range pathRe.FindAllStringIndex(msg, -1) {
|
||||
start, end := idx[0], idx[1]
|
||||
if start > 0 && isWordChar(msg[start-1]) {
|
||||
continue
|
||||
}
|
||||
b.WriteString(msg[last:start])
|
||||
b.WriteString("<redacted>")
|
||||
last = end
|
||||
}
|
||||
b.WriteString(msg[last:])
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func isWordChar(c byte) bool {
|
||||
return c == '_' || (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9')
|
||||
}
|
||||
|
||||
func truncate(b []byte) string {
|
||||
const limit = 200
|
||||
if len(b) <= limit {
|
||||
return string(b)
|
||||
}
|
||||
return string(b[:limit]) + "..."
|
||||
}
|
||||
@@ -0,0 +1,203 @@
|
||||
package gomsg
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// startRPCServer 启动带两路由的 RPC 服务端。
|
||||
func startRPCServer(t *testing.T) (*RpcServer, string) {
|
||||
t.Helper()
|
||||
s := NewRpcServer("127.0.0.1:0")
|
||||
s.AddRoute("echo", func(req map[string]any) (any, error) {
|
||||
return req, nil
|
||||
})
|
||||
s.AddRoute("sum", func(req map[string]any) (any, error) {
|
||||
a := req["a"].(json.Number)
|
||||
b := req["b"].(json.Number)
|
||||
ai, _ := a.Int64()
|
||||
bi, _ := b.Int64()
|
||||
return ai + bi, nil
|
||||
})
|
||||
s.AddRoute("fail", func(req map[string]any) (any, error) {
|
||||
return nil, errors.New("boom: /opt/micro_servers/db2/xxx.py failed")
|
||||
})
|
||||
s.AddRoute("panic", func(req map[string]any) (any, error) {
|
||||
panic("panic in handler")
|
||||
})
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(s.Stop)
|
||||
return s, s.Addr().String()
|
||||
}
|
||||
|
||||
func TestRpcClientCallSuccess(t *testing.T) {
|
||||
_, addr := startRPCServer(t)
|
||||
c := NewRpcClient(addr, 0)
|
||||
defer c.Close()
|
||||
|
||||
data, err := c.Call("echo", map[string]any{"hello": "world"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
m := data.(map[string]any)
|
||||
if m["hello"] != "world" {
|
||||
t.Fatalf("data = %v", data)
|
||||
}
|
||||
if _, has := m["_route"]; has {
|
||||
t.Fatal("_route should be stripped before handler")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRpcClientNumberPrecision(t *testing.T) {
|
||||
_, addr := startRPCServer(t)
|
||||
c := NewRpcClient(addr, 0)
|
||||
defer c.Close()
|
||||
|
||||
data, err := c.Call("sum", map[string]any{"a": json.Number("9007199254740993"), "b": json.Number("1")})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := data.(json.Number).String(); got != "9007199254740994" {
|
||||
t.Fatalf("sum = %s, want exact int64 precision", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRpcClientBusinessError(t *testing.T) {
|
||||
_, addr := startRPCServer(t)
|
||||
c := NewRpcClient(addr, 0)
|
||||
defer c.Close()
|
||||
|
||||
_, err := c.Call("fail", nil)
|
||||
var rpcErr *RpcError
|
||||
if !errors.As(err, &rpcErr) {
|
||||
t.Fatalf("expected *RpcError, got %T: %v", err, err)
|
||||
}
|
||||
// 绝对路径应被脱敏
|
||||
if strings.Contains(rpcErr.Error(), "/opt/") {
|
||||
t.Fatalf("error not sanitized: %v", rpcErr)
|
||||
}
|
||||
if !strings.Contains(rpcErr.Error(), "<redacted>") {
|
||||
t.Fatalf("expected <redacted> in error: %v", rpcErr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRpcUnknownRoute(t *testing.T) {
|
||||
_, addr := startRPCServer(t)
|
||||
c := NewRpcClient(addr, 0)
|
||||
defer c.Close()
|
||||
|
||||
_, err := c.Call("no_such_route", nil)
|
||||
if err == nil || !strings.Contains(err.Error(), "未知路由") {
|
||||
t.Fatalf("expected 未知路由 error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRpcServerInvalidJSON(t *testing.T) {
|
||||
_, addr := startRPCServer(t)
|
||||
c := NewMsgClient(addr, 0)
|
||||
defer c.Close()
|
||||
|
||||
resp, err := c.Call([]byte("{not json"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var r map[string]any
|
||||
if json.Unmarshal(resp, &r) != nil {
|
||||
t.Fatalf("invalid response: %q", resp)
|
||||
}
|
||||
if r["success"] != false || !strings.Contains(r["error"].(string), "invalid JSON request") {
|
||||
t.Fatalf("unexpected response: %v", r)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRpcServerPanicHandler(t *testing.T) {
|
||||
_, addr := startRPCServer(t)
|
||||
c := NewRpcClient(addr, 0)
|
||||
defer c.Close()
|
||||
|
||||
// handler 内 panic 不应打崩服务,后续请求仍可用
|
||||
if _, err := c.Call("panic", nil); err == nil {
|
||||
t.Fatal("expected error from panicking handler")
|
||||
}
|
||||
if _, err := c.Call("echo", map[string]any{"ok": true}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRpcServerOnErrorHook(t *testing.T) {
|
||||
s := NewRpcServer("127.0.0.1:0")
|
||||
var hookRoute string
|
||||
s.OnError = func(route string, err error) { hookRoute = route }
|
||||
s.AddRoute("fail", func(req map[string]any) (any, error) {
|
||||
return nil, errors.New("boom")
|
||||
})
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Stop()
|
||||
|
||||
c := NewRpcClient(s.Addr().String(), 0)
|
||||
defer c.Close()
|
||||
_, _ = c.Call("fail", nil)
|
||||
if hookRoute != "fail" {
|
||||
t.Fatalf("OnError hook route = %q", hookRoute)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRpcServerNonDictRequest(t *testing.T) {
|
||||
s := NewRpcServer("127.0.0.1:0")
|
||||
s.AddRoute("echo", func(req map[string]any) (any, error) { return req, nil })
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Stop()
|
||||
|
||||
c := NewMsgClient(s.Addr().String(), 0)
|
||||
defer c.Close()
|
||||
resp, err := c.Call([]byte(`["not","a","dict"]`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var r map[string]any
|
||||
if json.Unmarshal(resp, &r) != nil || r["success"] != false {
|
||||
t.Fatalf("expected failure response, got %q", resp)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSanitizeError(t *testing.T) {
|
||||
cases := []struct{ in, want string }{
|
||||
{"", ""},
|
||||
{"boom", "boom"},
|
||||
{"failed: /opt/micro_servers/db2/xxx.py", "failed: <redacted>"},
|
||||
{"no abs path here", "no abs path here"},
|
||||
{"abc/def 保持原样", "abc/def 保持原样"}, // 前导字符为单词字符时不匹配(对齐 Python lookbehind)
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := sanitizeError(c.in); got != c.want {
|
||||
t.Errorf("sanitizeError(%q) = %q, want %q", c.in, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRpcClientTimeout(t *testing.T) {
|
||||
s := NewRpcServer("127.0.0.1:0")
|
||||
s.AddRoute("slow", func(req map[string]any) (any, error) {
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
return nil, nil
|
||||
})
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Stop()
|
||||
|
||||
c := NewRpcClient(s.Addr().String(), 50*time.Millisecond)
|
||||
defer c.Close()
|
||||
if _, err := c.Call("slow", nil); err == nil {
|
||||
t.Fatal("expected timeout error")
|
||||
}
|
||||
}
|
||||
+130
@@ -0,0 +1,130 @@
|
||||
package gomsg
|
||||
|
||||
import (
|
||||
"net"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// MsgHandler 处理一条完整消息,返回回复载荷;返回 nil 表示不回复。
|
||||
type MsgHandler func(payload []byte) []byte
|
||||
|
||||
// MsgServer pymsg 协议服务端:accept 循环 + 每连接一个 goroutine 串行处理。
|
||||
// 对应 Python pymsg.MsgServer / C++ libmsg.MsgServer。
|
||||
type MsgServer struct {
|
||||
addr string
|
||||
handler MsgHandler
|
||||
// ReadTimeout 单条消息读取总超时;0 表示不设超时(默认,与 Python 行为一致)。
|
||||
ReadTimeout time.Duration
|
||||
|
||||
mu sync.Mutex
|
||||
ln net.Listener
|
||||
conns map[net.Conn]struct{}
|
||||
closed bool
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
// NewMsgServer 创建服务端(不监听,需调用 Start)。
|
||||
func NewMsgServer(addr string, handler MsgHandler) *MsgServer {
|
||||
return &MsgServer{
|
||||
addr: addr,
|
||||
handler: handler,
|
||||
conns: make(map[net.Conn]struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Addr 返回监听地址(Start 后可用,便于测试取随机端口)。
|
||||
func (s *MsgServer) Addr() net.Addr {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.ln == nil {
|
||||
return nil
|
||||
}
|
||||
return s.ln.Addr()
|
||||
}
|
||||
|
||||
// Start 开始监听并启动 accept 循环(非阻塞)。
|
||||
func (s *MsgServer) Start() error {
|
||||
ln, err := net.Listen("tcp", s.addr)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s.mu.Lock()
|
||||
if s.closed {
|
||||
s.mu.Unlock()
|
||||
_ = ln.Close()
|
||||
return nil
|
||||
}
|
||||
s.ln = ln
|
||||
s.mu.Unlock()
|
||||
|
||||
s.wg.Add(1)
|
||||
go s.acceptLoop(ln)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Stop 关闭监听与全部连接,等待处理中的请求结束。
|
||||
func (s *MsgServer) Stop() {
|
||||
s.mu.Lock()
|
||||
if s.closed {
|
||||
s.mu.Unlock()
|
||||
return
|
||||
}
|
||||
s.closed = true
|
||||
if s.ln != nil {
|
||||
_ = s.ln.Close()
|
||||
}
|
||||
for conn := range s.conns {
|
||||
_ = conn.Close()
|
||||
}
|
||||
s.mu.Unlock()
|
||||
s.wg.Wait()
|
||||
}
|
||||
|
||||
func (s *MsgServer) acceptLoop(ln net.Listener) {
|
||||
defer s.wg.Done()
|
||||
for {
|
||||
conn, err := ln.Accept()
|
||||
if err != nil {
|
||||
return // 监听关闭
|
||||
}
|
||||
s.mu.Lock()
|
||||
if s.closed {
|
||||
s.mu.Unlock()
|
||||
_ = conn.Close()
|
||||
return
|
||||
}
|
||||
s.conns[conn] = struct{}{}
|
||||
s.mu.Unlock()
|
||||
|
||||
s.wg.Add(1)
|
||||
go s.handleConn(conn)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *MsgServer) handleConn(conn net.Conn) {
|
||||
defer s.wg.Done()
|
||||
defer func() {
|
||||
s.mu.Lock()
|
||||
delete(s.conns, conn)
|
||||
s.mu.Unlock()
|
||||
_ = conn.Close()
|
||||
}()
|
||||
|
||||
for {
|
||||
if s.ReadTimeout > 0 {
|
||||
_ = conn.SetReadDeadline(time.Now().Add(s.ReadTimeout))
|
||||
}
|
||||
payload, err := ReadFrame(conn, 0)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
reply := s.handler(payload)
|
||||
if reply == nil {
|
||||
continue
|
||||
}
|
||||
if err := WriteFrame(conn, reply); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
robot_agent
|
||||
@@ -0,0 +1,22 @@
|
||||
# robot_agent — 基于 libhttp 的 DeepSeek 视觉机器人 agent 骨架
|
||||
CXX := g++
|
||||
CPPFLAGS := --std=c++20 -O2 -Wall -Wextra
|
||||
|
||||
LIBHTTP_DIR := ../libhttp
|
||||
INCS := -I $(LIBHTTP_DIR)/inc
|
||||
LIBS := -L $(LIBHTTP_DIR)/build -lus_httpclient -lboost_json -lssl -lcrypto -lcurl -ljpeg -lpthread
|
||||
RPATH := -Wl,-rpath,$(LIBHTTP_DIR)/build
|
||||
|
||||
TARGET := robot_agent
|
||||
|
||||
SRCS := main.cpp agent.cpp subagent.cpp tools.cpp camera.cpp
|
||||
|
||||
all: $(TARGET)
|
||||
|
||||
$(TARGET): $(SRCS) tools.hpp
|
||||
$(CXX) $(CPPFLAGS) $(INCS) $(SRCS) $(LIBS) $(RPATH) -o $@
|
||||
|
||||
clean:
|
||||
rm -f $(TARGET)
|
||||
|
||||
.PHONY: all clean
|
||||
@@ -0,0 +1,78 @@
|
||||
# robot_agent — 基于 libhttp 的 DeepSeek 视觉机器人 Agent
|
||||
|
||||
C++20 机器人智能体骨架:命令行输入任务,调用 DeepSeek 视觉大模型(`deepseek-v4-flash-vision-exp`),通过**工具循环**驱动摄像头识图、状态查询、命令行执行与文件读写,后续可扩展硬件控制。
|
||||
|
||||
## 架构
|
||||
|
||||
```
|
||||
用户任务(CLI 参数或 --prompt-file 文件)
|
||||
└─ 主代理(run_chat_loop 工具循环)
|
||||
├─ capture_image 拍照识图(真实摄像头 / 模拟图片)
|
||||
├─ get_robot_status 查询机器人状态
|
||||
└─ spawn_subagent 委派任务给子代理
|
||||
└─ 子代理(技能超集,独立工具循环)
|
||||
├─ run_command 执行 shell 命令(超时 + 输出截断)
|
||||
├─ read_file / write_file
|
||||
├─ capture_image / get_robot_status
|
||||
└─ 返回结论 → 主代理综合回答
|
||||
```
|
||||
|
||||
- **工具循环**:模型返回 `tool_calls` → agent 执行工具 → 结果回传 → 继续,直到模型给出最终回答(`--max-turns` 限轮)。
|
||||
- **图片注入**:OpenAI 协议的 tool 消息只允许文本,工具返回的图片由主循环作为 `image_url` 注入下一轮。
|
||||
- **时间戳**:会话开始时向 system prompt 注入一次当前时间(会话内不变);每个工具结果附 `工具执行时间`,供模型判断数据新旧。
|
||||
- **摄像头**:V4L2 直采,优先 MJPEG(硬件编码零转码,最快),回退 YUYV + libjpeg 编码;设备常开复用;`/dev/video*` 自动探测。
|
||||
|
||||
## 依赖与构建
|
||||
|
||||
- g++(C++20)、libhttp(`libus_httpclient.so`)、Boost.JSON、libcurl、OpenSSL、libjpeg、pthread
|
||||
- libhttp 需先编译:`cd ../libhttp && make`
|
||||
|
||||
```bash
|
||||
make # 生成 ./robot_agent
|
||||
make clean
|
||||
```
|
||||
|
||||
## 使用
|
||||
|
||||
```bash
|
||||
export DEEPSEEK_API_KEY=sk-xxx # 或用 --key
|
||||
|
||||
./robot_agent "看看环境里有什么" # 真实摄像头识图
|
||||
./robot_agent "报告状态" # 状态查询
|
||||
./robot_agent "列出当前目录文件" # 自动委派子代理执行命令
|
||||
./robot_agent --prompt-file /tmp/task.txt # 从文件读取任务
|
||||
./robot_agent "描述图片" --camera ./test.png # 模拟摄像头(读文件,便于离线测试)
|
||||
```
|
||||
|
||||
| 参数 | 说明 | 默认 |
|
||||
|---|---|---|
|
||||
| `--prompt-file <文件>` | 从文件读取任务(与位置参数互斥) | — |
|
||||
| `--camera <文件>` | 模拟摄像头:拍照工具读该图片 | 空 = 真实摄像头 |
|
||||
| `--camera-device <设备>` | 摄像头设备 | 自动探测 `/dev/video*` |
|
||||
| `--model <名称>` | 模型名 | `deepseek-v4-flash-vision-exp` |
|
||||
| `--key <API_KEY>` | API key | 环境变量 `DEEPSEEK_API_KEY` |
|
||||
| `--endpoint <URL>` | 接口地址 | `https://api.deepseek.com/chat/completions` |
|
||||
| `--max-turns <N>` | 最大工具轮次 | 8 |
|
||||
|
||||
工具日志输出到 stderr(`[主代理][工具]` / `[子代理][工具]`),stdout 只输出最终回答,便于脚本管道。
|
||||
|
||||
## 代码结构
|
||||
|
||||
| 文件 | 职责 |
|
||||
|---|---|
|
||||
| `main.cpp` | CLI 解析、主代理工具注册 |
|
||||
| `agent.hpp/cpp` | `run_chat_loop` 工具循环(主/子代理共用)、消息构造、时间戳注入 |
|
||||
| `subagent.hpp/cpp` | 子代理:技能超集 + 独立提示词 |
|
||||
| `tools.hpp/cpp` | `Tool`/`ToolRegistry` 与全部工具实现 |
|
||||
| `camera.hpp/cpp` | V4L2 抓帧(MJPEG/YUYV、mmap、超时保护) |
|
||||
|
||||
## 扩展
|
||||
|
||||
- **加工具**:实现 `Tool{name, description, parameters(JSON Schema), execute}` 后 `registry.add(...)`。
|
||||
- **工具返回图片**:`ToolResult.image_data_url` 填 data URL,主循环自动注入下一轮。
|
||||
- **加子代理类型**:仿照 `subagent.cpp` 组装不同工具集与提示词。
|
||||
- **接真实硬件**:把 `capture_image`/`get_robot_status` 的实现换成设备驱动调用即可,接口不变。
|
||||
|
||||
## 安全注意
|
||||
|
||||
`run_command` 会执行任意 shell 命令且无白名单/审批,仅建议在可信环境(你自己的机器人主机)运行;后续如需对外暴露,应先加命令审批或沙箱。
|
||||
@@ -0,0 +1,155 @@
|
||||
#include "agent.hpp"
|
||||
|
||||
#include <us/HttpClient.hpp>
|
||||
|
||||
#include <boost/system/error_code.hpp>
|
||||
|
||||
#include <ctime>
|
||||
#include <iostream>
|
||||
#include <stdexcept>
|
||||
|
||||
using pmc::net::HttpClient;
|
||||
|
||||
namespace {
|
||||
|
||||
// 本地时间 ISO8601,如 2026-08-21T14:30:05+08:00
|
||||
std::string now_iso8601() {
|
||||
std::time_t tt = std::chrono::system_clock::to_time_t(std::chrono::system_clock::now());
|
||||
std::tm tm{};
|
||||
localtime_r(&tt, &tm);
|
||||
char buf[64];
|
||||
strftime(buf, sizeof(buf), "%Y-%m-%dT%H:%M:%S%z", &tm);
|
||||
std::string s(buf);
|
||||
if (s.size() >= 5 && (s[s.size() - 5] == '+' || s[s.size() - 5] == '-'))
|
||||
s.insert(s.size() - 2, ":");
|
||||
return s;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
json::object user_text_msg(const std::string& text) {
|
||||
return json::object{{"role", "user"}, {"content", text}};
|
||||
}
|
||||
|
||||
json::object user_image_msg(const std::string& text, const std::string& data_url) {
|
||||
return json::object{
|
||||
{"role", "user"},
|
||||
{"content", json::array{
|
||||
json::object{{"type", "text"}, {"text", text}},
|
||||
json::object{{"type", "image_url"},
|
||||
{"image_url", json::object{{"url", data_url}}}},
|
||||
}},
|
||||
};
|
||||
}
|
||||
|
||||
json::object tool_result_msg(const std::string& call_id, const std::string& text) {
|
||||
return json::object{{"role", "tool"}, {"tool_call_id", call_id}, {"content", text}};
|
||||
}
|
||||
|
||||
json::value build_tools_spec(const ToolRegistry& registry) {
|
||||
json::array arr;
|
||||
for (const auto& t : registry.all()) {
|
||||
arr.push_back(json::object{
|
||||
{"type", "function"},
|
||||
{"function", json::object{
|
||||
{"name", t.name},
|
||||
{"description", t.description},
|
||||
{"parameters", t.parameters},
|
||||
}},
|
||||
});
|
||||
}
|
||||
return arr;
|
||||
}
|
||||
|
||||
std::string run_chat_loop(const AgentConfig& cfg,
|
||||
const std::string& system_prompt,
|
||||
const std::string& user_question,
|
||||
const ToolRegistry& registry,
|
||||
const std::string& log_tag) {
|
||||
json::array messages;
|
||||
// 会话开始注入一次当前时间(会话内保持不变,避免模型混淆"现在")
|
||||
std::string timed_prompt =
|
||||
system_prompt + "\n当前时间: " + now_iso8601() + "(本会话内保持不变)";
|
||||
messages.push_back(json::object{{"role", "system"}, {"content", timed_prompt}});
|
||||
messages.push_back(user_text_msg(user_question));
|
||||
|
||||
HttpClient client;
|
||||
client.set_header("Authorization: Bearer " + cfg.api_key);
|
||||
client.set_header("Content-Type: application/json");
|
||||
|
||||
for (int turn = 0; turn < cfg.max_turns; ++turn) {
|
||||
json::object body;
|
||||
body["model"] = cfg.model;
|
||||
body["messages"] = messages;
|
||||
body["tools"] = build_tools_spec(registry);
|
||||
body["stream"] = false;
|
||||
|
||||
std::string response_body;
|
||||
client.post_stream_sync(cfg.endpoint, json::serialize(body),
|
||||
[&](const std::string& chunk) { response_body += chunk; });
|
||||
|
||||
long code = client.get_response_code();
|
||||
if (code != 200)
|
||||
throw std::runtime_error("HTTP " + std::to_string(code) + ": " + response_body);
|
||||
|
||||
boost::system::error_code ec;
|
||||
json::value resp = json::parse(response_body, ec);
|
||||
if (ec)
|
||||
throw std::runtime_error("响应不是合法 JSON: " + response_body.substr(0, 200));
|
||||
if (resp.as_object().contains("error"))
|
||||
throw std::runtime_error("API 错误: " + json::serialize(resp.at("error")));
|
||||
|
||||
auto& choices = resp.at("choices").as_array();
|
||||
if (choices.empty()) throw std::runtime_error("响应中没有 choices");
|
||||
auto& message = choices[0].as_object().at("message").as_object();
|
||||
|
||||
auto tc_it = message.find("tool_calls");
|
||||
bool has_tool_calls = tc_it != message.end() && !tc_it->value().is_null() &&
|
||||
!tc_it->value().as_array().empty();
|
||||
|
||||
if (has_tool_calls) {
|
||||
messages.push_back(message); // 回传 assistant 的 tool_calls 消息
|
||||
|
||||
for (const auto& call : tc_it->value().as_array()) {
|
||||
const auto& c = call.as_object();
|
||||
std::string call_id =
|
||||
c.contains("id") ? json::value_to<std::string>(c.at("id")) : "";
|
||||
const auto& fn = c.at("function").as_object();
|
||||
std::string name = json::value_to<std::string>(fn.at("name"));
|
||||
|
||||
json::value args;
|
||||
auto args_it = fn.find("arguments");
|
||||
if (args_it != fn.end() && !args_it->value().is_null()) {
|
||||
std::string args_str = json::value_to<std::string>(args_it->value());
|
||||
boost::system::error_code aec;
|
||||
args = json::parse(args_str, aec);
|
||||
if (aec) args = json::object();
|
||||
} else {
|
||||
args = json::object();
|
||||
}
|
||||
|
||||
ToolResult result;
|
||||
if (!registry.execute(name, args, result))
|
||||
result = {"错误: 未知工具 " + name, ""};
|
||||
|
||||
// 工具结果统一附执行时间戳,供模型判断数据新旧
|
||||
std::string timed_result = "工具执行时间: " + now_iso8601() + "\n" + result.text;
|
||||
std::cerr << "[" << log_tag << "][工具] " << name << " → "
|
||||
<< timed_result.substr(0, 120) << "\n";
|
||||
messages.push_back(tool_result_msg(call_id, timed_result));
|
||||
if (!result.image_data_url.empty())
|
||||
messages.push_back(
|
||||
user_image_msg("工具返回的图片如下,请结合图片继续回答。",
|
||||
result.image_data_url));
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
if (message.contains("content") && message.at("content").is_string())
|
||||
return json::value_to<std::string>(message.at("content"));
|
||||
throw std::runtime_error("模型没有返回文本内容");
|
||||
}
|
||||
|
||||
throw std::runtime_error("达到最大工具轮次 (" + std::to_string(cfg.max_turns) +
|
||||
"),未获得最终回答");
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
#pragma once
|
||||
|
||||
#include <boost/json.hpp>
|
||||
|
||||
#include "tools.hpp"
|
||||
|
||||
#include <string>
|
||||
|
||||
namespace json = boost::json;
|
||||
|
||||
struct AgentConfig {
|
||||
std::string model;
|
||||
std::string endpoint;
|
||||
std::string api_key;
|
||||
std::string camera_path; // 非空 = 模拟摄像头(读文件)
|
||||
std::string camera_device; // 真实摄像头设备(默认 /dev/video0)
|
||||
int max_turns = 8;
|
||||
};
|
||||
|
||||
// 运行工具循环,直到模型给出最终回答并返回其文本。失败抛 std::runtime_error。
|
||||
// log_tag 用于区分不同代理的工具日志(输出到 stderr)。
|
||||
std::string run_chat_loop(const AgentConfig& cfg,
|
||||
const std::string& system_prompt,
|
||||
const std::string& user_question,
|
||||
const ToolRegistry& registry,
|
||||
const std::string& log_tag);
|
||||
|
||||
// 消息构造(主代理与子代理共用)
|
||||
json::object user_text_msg(const std::string& text);
|
||||
json::object user_image_msg(const std::string& text, const std::string& data_url);
|
||||
json::object tool_result_msg(const std::string& call_id, const std::string& text);
|
||||
json::value build_tools_spec(const ToolRegistry& registry);
|
||||
@@ -0,0 +1,249 @@
|
||||
#include "camera.hpp"
|
||||
|
||||
#include <jpeglib.h>
|
||||
|
||||
#include <fcntl.h>
|
||||
#include <glob.h>
|
||||
#include <linux/videodev2.h>
|
||||
#include <poll.h>
|
||||
#include <sys/ioctl.h>
|
||||
#include <sys/mman.h>
|
||||
#include <unistd.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <chrono>
|
||||
#include <cstring>
|
||||
#include <stdexcept>
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr size_t CAPTURE_BUFFER_COUNT = 3;
|
||||
constexpr int FRAME_TIMEOUT_MS = 1500; // 单帧出队超时
|
||||
|
||||
// 扫描 /dev/video*,返回第一个支持视频采集的设备
|
||||
std::string auto_detect_device() {
|
||||
glob_t g;
|
||||
std::vector<std::string> candidates;
|
||||
if (glob("/dev/video*", 0, nullptr, &g) == 0) {
|
||||
for (size_t i = 0; i < g.gl_pathc; ++i)
|
||||
candidates.emplace_back(g.gl_pathv[i]);
|
||||
}
|
||||
globfree(&g);
|
||||
std::sort(candidates.begin(), candidates.end());
|
||||
|
||||
for (const auto& dev : candidates) {
|
||||
int fd = open(dev.c_str(), O_RDWR | O_NONBLOCK);
|
||||
if (fd < 0) continue;
|
||||
v4l2_capability cap{};
|
||||
bool ok = ioctl(fd, VIDIOC_QUERYCAP, &cap) == 0 &&
|
||||
(cap.capabilities & V4L2_CAP_VIDEO_CAPTURE);
|
||||
close(fd);
|
||||
if (ok) return dev;
|
||||
}
|
||||
throw std::runtime_error("未找到可用的视频采集设备(已扫描 /dev/video*,"
|
||||
"可用 --camera-device 指定)");
|
||||
}
|
||||
|
||||
void xioctl(int fd, unsigned long request, void* arg, const char* what) {
|
||||
if (ioctl(fd, request, arg) != 0)
|
||||
throw std::runtime_error(std::string(what) + " 失败: " + strerror(errno));
|
||||
}
|
||||
|
||||
// YUYV → RGB24 → JPEG(libjpeg 内存压缩)
|
||||
std::vector<unsigned char> yuyv_to_jpeg(const std::vector<unsigned char>& yuyv,
|
||||
unsigned w, unsigned h) {
|
||||
std::vector<unsigned char> rgb(static_cast<size_t>(w) * h * 3);
|
||||
const unsigned char* in = yuyv.data();
|
||||
unsigned char* out = rgb.data();
|
||||
for (unsigned y = 0; y < h; ++y) {
|
||||
for (unsigned x = 0; x + 1 < w; x += 2) {
|
||||
int y0 = in[0], u = in[1] - 128, y1 = in[2], v = in[3] - 128;
|
||||
in += 4;
|
||||
int c0 = y0 - 16, c1 = y1 - 16;
|
||||
int r0 = (298 * c0 + 409 * v + 128) >> 8;
|
||||
int g0 = (298 * c0 - 100 * u - 208 * v + 128) >> 8;
|
||||
int b0 = (298 * c0 + 516 * u + 128) >> 8;
|
||||
int r1 = (298 * c1 + 409 * v + 128) >> 8;
|
||||
int g1 = (298 * c1 - 100 * u - 208 * v + 128) >> 8;
|
||||
int b1 = (298 * c1 + 516 * u + 128) >> 8;
|
||||
*out++ = static_cast<unsigned char>(std::clamp(r0, 0, 255));
|
||||
*out++ = static_cast<unsigned char>(std::clamp(g0, 0, 255));
|
||||
*out++ = static_cast<unsigned char>(std::clamp(b0, 0, 255));
|
||||
*out++ = static_cast<unsigned char>(std::clamp(r1, 0, 255));
|
||||
*out++ = static_cast<unsigned char>(std::clamp(g1, 0, 255));
|
||||
*out++ = static_cast<unsigned char>(std::clamp(b1, 0, 255));
|
||||
}
|
||||
if (w % 2) { // 奇数宽度:最后一像素复用前一对的 U/V
|
||||
int y0 = in[0], u = in[1] - 128, v = in[2] - 128;
|
||||
in += 4;
|
||||
int c0 = y0 - 16;
|
||||
*out++ = static_cast<unsigned char>(std::clamp((298 * c0 + 409 * v + 128) >> 8, 0, 255));
|
||||
*out++ = static_cast<unsigned char>(std::clamp((298 * c0 - 100 * u - 208 * v + 128) >> 8, 0, 255));
|
||||
*out++ = static_cast<unsigned char>(std::clamp((298 * c0 + 516 * u + 128) >> 8, 0, 255));
|
||||
}
|
||||
}
|
||||
|
||||
jpeg_compress_struct cinfo;
|
||||
jpeg_error_mgr jerr;
|
||||
cinfo.err = jpeg_std_error(&jerr);
|
||||
jpeg_create_compress(&cinfo);
|
||||
unsigned char* outbuf = nullptr;
|
||||
unsigned long outsize = 0;
|
||||
jpeg_mem_dest(&cinfo, &outbuf, &outsize);
|
||||
cinfo.image_width = w;
|
||||
cinfo.image_height = h;
|
||||
cinfo.input_components = 3;
|
||||
cinfo.in_color_space = JCS_RGB;
|
||||
jpeg_set_defaults(&cinfo);
|
||||
jpeg_set_quality(&cinfo, 85, TRUE);
|
||||
jpeg_start_compress(&cinfo, TRUE);
|
||||
while (cinfo.next_scanline < h) {
|
||||
JSAMPROW row = rgb.data() + static_cast<size_t>(cinfo.next_scanline) * w * 3;
|
||||
jpeg_write_scanlines(&cinfo, &row, 1);
|
||||
}
|
||||
jpeg_finish_compress(&cinfo);
|
||||
std::vector<unsigned char> result(outbuf, outbuf + outsize);
|
||||
free(outbuf);
|
||||
jpeg_destroy_compress(&cinfo);
|
||||
return result;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
CameraCapture::CameraCapture(const std::string& device) : device_(device) {
|
||||
if (device_.empty()) device_ = auto_detect_device();
|
||||
open_device();
|
||||
init_format();
|
||||
request_buffers();
|
||||
}
|
||||
|
||||
CameraCapture::~CameraCapture() {
|
||||
std::lock_guard<std::mutex> lock(mtx_);
|
||||
if (fd_ < 0) return;
|
||||
try {
|
||||
stream_off();
|
||||
} catch (...) {
|
||||
}
|
||||
for (size_t i = 0; i < buf_count_; ++i)
|
||||
munmap(buffers_[i], buffer_sizes_[i]);
|
||||
close(fd_);
|
||||
}
|
||||
|
||||
void CameraCapture::open_device() {
|
||||
fd_ = open(device_.c_str(), O_RDWR | O_NONBLOCK);
|
||||
if (fd_ < 0)
|
||||
throw std::runtime_error("无法打开摄像头 " + device_ + ": " + strerror(errno));
|
||||
|
||||
v4l2_capability cap{};
|
||||
if (ioctl(fd_, VIDIOC_QUERYCAP, &cap) != 0)
|
||||
throw std::runtime_error("VIDIOC_QUERYCAP 失败: " + std::string(strerror(errno)));
|
||||
if (!(cap.capabilities & V4L2_CAP_VIDEO_CAPTURE) ||
|
||||
!(cap.capabilities & V4L2_CAP_STREAMING))
|
||||
throw std::runtime_error(device_ + " 不支持视频捕获或 mmap 流式传输");
|
||||
}
|
||||
|
||||
void CameraCapture::init_format() {
|
||||
v4l2_format fmt{};
|
||||
fmt.type = V4L2_BUF_TYPE_VIDEO_CAPTURE;
|
||||
fmt.fmt.pix.width = 640;
|
||||
fmt.fmt.pix.height = 480;
|
||||
fmt.fmt.pix.field = V4L2_FIELD_ANY;
|
||||
|
||||
// 优先 MJPEG:摄像头硬件编码,抓帧即用,无需转码
|
||||
fmt.fmt.pix.pixelformat = V4L2_PIX_FMT_MJPEG;
|
||||
if (ioctl(fd_, VIDIOC_S_FMT, &fmt) == 0 && fmt.fmt.pix.pixelformat == V4L2_PIX_FMT_MJPEG) {
|
||||
jpeg_format_ = true;
|
||||
} else {
|
||||
fmt.fmt.pix.pixelformat = V4L2_PIX_FMT_YUYV;
|
||||
xioctl(fd_, VIDIOC_S_FMT, &fmt, "设置 YUYV 格式");
|
||||
if (fmt.fmt.pix.pixelformat != V4L2_PIX_FMT_YUYV)
|
||||
throw std::runtime_error(device_ + " 不支持 MJPEG 或 YUYV");
|
||||
jpeg_format_ = false;
|
||||
}
|
||||
width_ = fmt.fmt.pix.width;
|
||||
height_ = fmt.fmt.pix.height;
|
||||
}
|
||||
|
||||
void CameraCapture::request_buffers() {
|
||||
v4l2_requestbuffers req{};
|
||||
req.count = CAPTURE_BUFFER_COUNT;
|
||||
req.type = V4L2_BUF_TYPE_VIDEO_CAPTURE;
|
||||
req.memory = V4L2_MEMORY_MMAP;
|
||||
xioctl(fd_, VIDIOC_REQBUFS, &req, "申请缓冲");
|
||||
if (req.count < 2)
|
||||
throw std::runtime_error("摄像头缓冲不足");
|
||||
buf_count_ = req.count;
|
||||
|
||||
buffers_.resize(buf_count_);
|
||||
buffer_sizes_.resize(buf_count_);
|
||||
for (size_t i = 0; i < buf_count_; ++i) {
|
||||
v4l2_buffer buf{};
|
||||
buf.type = V4L2_BUF_TYPE_VIDEO_CAPTURE;
|
||||
buf.memory = V4L2_MEMORY_MMAP;
|
||||
buf.index = i;
|
||||
xioctl(fd_, VIDIOC_QUERYBUF, &buf, "查询缓冲");
|
||||
buffers_[i] = mmap(nullptr, buf.length, PROT_READ | PROT_WRITE, MAP_SHARED, fd_,
|
||||
buf.m.offset);
|
||||
if (buffers_[i] == MAP_FAILED)
|
||||
throw std::runtime_error("mmap 失败: " + std::string(strerror(errno)));
|
||||
buffer_sizes_[i] = buf.length;
|
||||
xioctl(fd_, VIDIOC_QBUF, &buf, "缓冲入队");
|
||||
}
|
||||
}
|
||||
|
||||
void CameraCapture::stream_on() {
|
||||
int type = V4L2_BUF_TYPE_VIDEO_CAPTURE;
|
||||
xioctl(fd_, VIDIOC_STREAMON, &type, "开启视频流");
|
||||
}
|
||||
|
||||
void CameraCapture::stream_off() {
|
||||
int type = V4L2_BUF_TYPE_VIDEO_CAPTURE;
|
||||
if (ioctl(fd_, VIDIOC_STREAMOFF, &type) != 0 && errno != EINVAL)
|
||||
throw std::runtime_error("关闭视频流失败: " + std::string(strerror(errno)));
|
||||
}
|
||||
|
||||
std::vector<unsigned char> CameraCapture::dequeue_frame(int timeout_ms) {
|
||||
auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(timeout_ms);
|
||||
while (true) {
|
||||
v4l2_buffer buf{};
|
||||
buf.type = V4L2_BUF_TYPE_VIDEO_CAPTURE;
|
||||
buf.memory = V4L2_MEMORY_MMAP;
|
||||
if (ioctl(fd_, VIDIOC_DQBUF, &buf) == 0) {
|
||||
std::vector<unsigned char> frame(static_cast<unsigned char*>(buffers_[buf.index]),
|
||||
static_cast<unsigned char*>(buffers_[buf.index]) +
|
||||
buf.bytesused);
|
||||
xioctl(fd_, VIDIOC_QBUF, &buf, "缓冲重新入队");
|
||||
return frame;
|
||||
}
|
||||
if (errno != EAGAIN)
|
||||
throw std::runtime_error("出队帧失败: " + std::string(strerror(errno)));
|
||||
auto now = std::chrono::steady_clock::now();
|
||||
if (now >= deadline)
|
||||
throw std::runtime_error("等待摄像头画面超时");
|
||||
int remain_ms = static_cast<int>(
|
||||
std::chrono::duration_cast<std::chrono::milliseconds>(deadline - now).count());
|
||||
struct pollfd pfd{fd_, POLLIN, 0};
|
||||
poll(&pfd, 1, std::min(remain_ms, 200));
|
||||
}
|
||||
}
|
||||
|
||||
std::string CameraCapture::grab_jpeg() {
|
||||
std::lock_guard<std::mutex> lock(mtx_);
|
||||
stream_on();
|
||||
try {
|
||||
dequeue_frame(FRAME_TIMEOUT_MS); // 丢弃第一帧(曝光未稳定)
|
||||
auto frame = dequeue_frame(FRAME_TIMEOUT_MS); // 取第二帧
|
||||
stream_off();
|
||||
if (!jpeg_format_) {
|
||||
auto jpeg = yuyv_to_jpeg(frame, width_, height_);
|
||||
return std::string(reinterpret_cast<char*>(jpeg.data()), jpeg.size());
|
||||
}
|
||||
return std::string(reinterpret_cast<char*>(frame.data()), frame.size());
|
||||
} catch (...) {
|
||||
try {
|
||||
stream_off();
|
||||
} catch (...) {
|
||||
}
|
||||
throw;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
#pragma once
|
||||
|
||||
#include <cstddef>
|
||||
#include <mutex>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
// V4L2 摄像头抓帧:优先 MJPEG(硬件编码,抓到的就是 JPEG,零转码,速度最快),
|
||||
// 设备不支持时回退 YUYV + libjpeg 软件编码。
|
||||
// 设备常开复用(构造时打开并建流),每次抓帧 = STREAMON → 取帧 → STREAMOFF。
|
||||
// device 为空时自动扫描 /dev/video* 探测第一个可用设备。失败抛 std::runtime_error。
|
||||
class CameraCapture {
|
||||
public:
|
||||
explicit CameraCapture(const std::string& device);
|
||||
~CameraCapture();
|
||||
|
||||
CameraCapture(const CameraCapture&) = delete;
|
||||
CameraCapture& operator=(const CameraCapture&) = delete;
|
||||
|
||||
// 抓一帧,返回 JPEG 数据
|
||||
std::string grab_jpeg();
|
||||
|
||||
unsigned width() const { return width_; }
|
||||
unsigned height() const { return height_; }
|
||||
|
||||
private:
|
||||
void open_device();
|
||||
void init_format();
|
||||
void request_buffers();
|
||||
void stream_on();
|
||||
void stream_off();
|
||||
// 出队一帧(带超时)。skip_first=true 时丢弃第一帧再取第二帧(等待曝光稳定)
|
||||
std::vector<unsigned char> dequeue_frame(int timeout_ms);
|
||||
|
||||
std::string device_;
|
||||
int fd_ = -1;
|
||||
bool jpeg_format_ = false; // true = MJPEG 直出,false = YUYV 需转码
|
||||
unsigned width_ = 0;
|
||||
unsigned height_ = 0;
|
||||
std::vector<void*> buffers_;
|
||||
std::vector<size_t> buffer_sizes_;
|
||||
size_t buf_count_ = 0;
|
||||
std::mutex mtx_;
|
||||
};
|
||||
@@ -0,0 +1,129 @@
|
||||
#include "agent.hpp"
|
||||
#include "tools.hpp"
|
||||
|
||||
#include <cstdlib>
|
||||
#include <fstream>
|
||||
#include <iostream>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr const char* DEFAULT_ENDPOINT = "https://api.deepseek.com/chat/completions";
|
||||
constexpr const char* DEFAULT_MODEL = "deepseek-v4-flash-vision-exp";
|
||||
|
||||
constexpr const char* MAIN_SYSTEM_PROMPT =
|
||||
"你是控制硬件机器人的主智能体。当用户的问题涉及观察环境、看东西、识别物体时,"
|
||||
"必须先调用 capture_image 工具获取照片再回答。"
|
||||
"当任务需要执行 shell 命令、读写文件等操作时,调用 spawn_subagent 工具委派给子代理,"
|
||||
"不要自己尝试执行。回答使用中文。";
|
||||
|
||||
struct Options {
|
||||
std::string question;
|
||||
std::string prompt_file; // 非空 = 从文件读取问题
|
||||
std::string camera_path; // 非空 = 模拟摄像头(读文件)
|
||||
std::string camera_device; // 空 = 自动探测 /dev/video*
|
||||
std::string model = DEFAULT_MODEL;
|
||||
std::string endpoint = DEFAULT_ENDPOINT;
|
||||
std::string api_key;
|
||||
int max_turns = 8;
|
||||
};
|
||||
|
||||
void print_usage(const char* prog) {
|
||||
std::cerr << "用法: " << prog << " \"问题\" [选项]\n"
|
||||
<< " " << prog << " --prompt-file <文件> [选项]\n"
|
||||
<< "选项:\n"
|
||||
<< " --prompt-file <文件> 从文件读取问题(与位置参数互斥)\n"
|
||||
<< " --camera <文件> 模拟摄像头:拍照工具读该图片(不指定则用真实摄像头)\n"
|
||||
<< " --camera-device <设备> 摄像头设备(默认自动探测 /dev/video*)\n"
|
||||
<< " --model <名称> 模型名(默认 " << DEFAULT_MODEL << ")\n"
|
||||
<< " --key <API_KEY> API key(默认读环境变量 DEEPSEEK_API_KEY)\n"
|
||||
<< " --endpoint <URL> 接口地址(默认 DeepSeek 官方)\n"
|
||||
<< " --max-turns <N> 最大工具轮次(默认 8)\n";
|
||||
}
|
||||
|
||||
Options parse_args(int argc, char** argv) {
|
||||
Options opt;
|
||||
if (const char* key_env = std::getenv("DEEPSEEK_API_KEY")) opt.api_key = key_env;
|
||||
|
||||
for (int i = 1; i < argc; ++i) {
|
||||
std::string arg = argv[i];
|
||||
auto next = [&](const char* name) -> std::string {
|
||||
if (i + 1 >= argc) {
|
||||
std::cerr << "错误: " << name << " 缺少参数\n";
|
||||
std::exit(2);
|
||||
}
|
||||
return argv[++i];
|
||||
};
|
||||
if (arg == "--camera") opt.camera_path = next("--camera");
|
||||
else if (arg == "--camera-device") opt.camera_device = next("--camera-device");
|
||||
else if (arg == "--prompt-file") opt.prompt_file = next("--prompt-file");
|
||||
else if (arg == "--model") opt.model = next("--model");
|
||||
else if (arg == "--key") opt.api_key = next("--key");
|
||||
else if (arg == "--endpoint") opt.endpoint = next("--endpoint");
|
||||
else if (arg == "--max-turns") opt.max_turns = std::stoi(next("--max-turns"));
|
||||
else if (arg == "--help" || arg == "-h") { print_usage(argv[0]); std::exit(0); }
|
||||
else if (!arg.empty() && arg[0] == '-') {
|
||||
std::cerr << "错误: 未知选项 " << arg << "\n";
|
||||
print_usage(argv[0]);
|
||||
std::exit(2);
|
||||
} else {
|
||||
opt.question = arg;
|
||||
}
|
||||
}
|
||||
return opt;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
Options opt = parse_args(argc, argv);
|
||||
|
||||
if (!opt.prompt_file.empty()) {
|
||||
if (!opt.question.empty()) {
|
||||
std::cerr << "错误: --prompt-file 与位置参数问题不能同时使用\n";
|
||||
return 2;
|
||||
}
|
||||
std::ifstream f(opt.prompt_file);
|
||||
if (!f) {
|
||||
std::cerr << "错误: 无法读取提示词文件 " << opt.prompt_file << "\n";
|
||||
return 2;
|
||||
}
|
||||
std::ostringstream ss;
|
||||
ss << f.rdbuf();
|
||||
opt.question = ss.str();
|
||||
}
|
||||
|
||||
if (opt.question.empty()) {
|
||||
std::cerr << "错误: 缺少问题参数(用位置参数或 --prompt-file)\n";
|
||||
print_usage(argv[0]);
|
||||
return 2;
|
||||
}
|
||||
if (opt.api_key.empty()) {
|
||||
std::cerr << "错误: 缺少 API key(用 --key 或设置环境变量 DEEPSEEK_API_KEY)\n";
|
||||
return 2;
|
||||
}
|
||||
|
||||
AgentConfig cfg;
|
||||
cfg.model = opt.model;
|
||||
cfg.endpoint = opt.endpoint;
|
||||
cfg.api_key = opt.api_key;
|
||||
cfg.camera_path = opt.camera_path;
|
||||
cfg.camera_device = opt.camera_device;
|
||||
cfg.max_turns = opt.max_turns;
|
||||
|
||||
ToolRegistry registry;
|
||||
registry.add(make_capture_image_tool(opt.camera_path, opt.camera_device));
|
||||
registry.add(make_robot_status_tool());
|
||||
registry.add(make_spawn_subagent_tool(cfg));
|
||||
|
||||
try {
|
||||
std::string answer =
|
||||
run_chat_loop(cfg, MAIN_SYSTEM_PROMPT, opt.question, registry, "主代理");
|
||||
std::cout << answer << "\n";
|
||||
} catch (const std::exception& e) {
|
||||
std::cerr << "错误: " << e.what() << "\n";
|
||||
return 1;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
#include "subagent.hpp"
|
||||
|
||||
#include "tools.hpp"
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr const char* SUBAGENT_PROMPT =
|
||||
"你是机器人的子代理(专家),拥有比主代理更丰富的技能:可以执行 shell 命令、"
|
||||
"读写文件、拍照和查询状态。请利用这些技能完成主代理委派的任务,"
|
||||
"返回简洁明确的结果,用中文。";
|
||||
|
||||
} // namespace
|
||||
|
||||
std::string run_subagent(const AgentConfig& cfg, const std::string& task) {
|
||||
ToolRegistry registry;
|
||||
registry.add(make_run_command_tool());
|
||||
registry.add(make_read_file_tool());
|
||||
registry.add(make_write_file_tool());
|
||||
registry.add(make_capture_image_tool(cfg.camera_path, cfg.camera_device));
|
||||
registry.add(make_robot_status_tool());
|
||||
|
||||
return run_chat_loop(cfg, SUBAGENT_PROMPT, task, registry, "子代理");
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
#pragma once
|
||||
|
||||
#include "agent.hpp"
|
||||
|
||||
#include <string>
|
||||
|
||||
// 子代理:拥有命令行执行、文件读写等扩展技能(主代理的超集),
|
||||
// 完成主代理委派的任务后返回结论文本。
|
||||
std::string run_subagent(const AgentConfig& cfg, const std::string& task);
|
||||
@@ -0,0 +1,5 @@
|
||||
# 身份
|
||||
你是一个机器人,你的操作系统应该是linux
|
||||
|
||||
# 任务
|
||||
使用摄像头观察情况,根据你看到的情况,通过命令行做出行动。
|
||||
@@ -0,0 +1,329 @@
|
||||
#include "tools.hpp"
|
||||
|
||||
#include "agent.hpp"
|
||||
#include "camera.hpp"
|
||||
#include "subagent.hpp"
|
||||
|
||||
#include <array>
|
||||
#include <cerrno>
|
||||
#include <chrono>
|
||||
#include <csignal>
|
||||
#include <fcntl.h>
|
||||
#include <fstream>
|
||||
#include <memory>
|
||||
#include <poll.h>
|
||||
#include <sstream>
|
||||
#include <stdexcept>
|
||||
#include <sys/wait.h>
|
||||
#include <unistd.h>
|
||||
|
||||
namespace {
|
||||
|
||||
std::string base64_encode(const std::string& in) {
|
||||
static const char table[] =
|
||||
"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
|
||||
std::string out;
|
||||
out.reserve((in.size() + 2) / 3 * 4);
|
||||
for (size_t i = 0; i < in.size(); i += 3) {
|
||||
unsigned n = static_cast<unsigned char>(in[i]) << 16;
|
||||
if (i + 1 < in.size()) n |= static_cast<unsigned char>(in[i + 1]) << 8;
|
||||
if (i + 2 < in.size()) n |= static_cast<unsigned char>(in[i + 2]);
|
||||
out.push_back(table[(n >> 18) & 0x3F]);
|
||||
out.push_back(table[(n >> 12) & 0x3F]);
|
||||
out.push_back(i + 1 < in.size() ? table[(n >> 6) & 0x3F] : '=');
|
||||
out.push_back(i + 2 < in.size() ? table[n & 0x3F] : '=');
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
std::string mime_from_path(const std::string& path) {
|
||||
auto pos = path.find_last_of('.');
|
||||
std::string ext = (pos == std::string::npos) ? "" : path.substr(pos + 1);
|
||||
if (ext == "jpg" || ext == "jpeg") return "image/jpeg";
|
||||
if (ext == "png") return "image/png";
|
||||
if (ext == "webp") return "image/webp";
|
||||
if (ext == "gif") return "image/gif";
|
||||
throw std::runtime_error("不支持的图片格式: " + ext);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
Tool make_capture_image_tool(const std::string& simulated_image_path,
|
||||
const std::string& camera_device) {
|
||||
Tool t;
|
||||
t.name = "capture_image";
|
||||
t.description =
|
||||
"拍摄一张机器人当前环境的照片。当用户的问题涉及观察、看、识别物体或了解周围环境时,"
|
||||
"必须先调用本工具获取照片,再基于照片回答。";
|
||||
t.parameters = json::object{{"type", "object"}, {"properties", json::object{}}};
|
||||
|
||||
// 摄像头实例懒加载并常开复用,避免每次抓帧都重新打开设备
|
||||
auto camera = std::make_shared<std::shared_ptr<CameraCapture>>();
|
||||
|
||||
t.execute = [simulated_image_path, camera_device, camera](const json::value&) -> ToolResult {
|
||||
if (!simulated_image_path.empty()) {
|
||||
std::ifstream f(simulated_image_path, std::ios::binary);
|
||||
if (!f)
|
||||
return {"错误: 无法读取模拟摄像头图片 " + simulated_image_path, ""};
|
||||
std::ostringstream ss;
|
||||
ss << f.rdbuf();
|
||||
std::string data_url;
|
||||
try {
|
||||
data_url = "data:" + mime_from_path(simulated_image_path) + ";base64," +
|
||||
base64_encode(ss.str());
|
||||
} catch (const std::exception& e) {
|
||||
return {std::string("错误: ") + e.what(), ""};
|
||||
}
|
||||
return {"已拍摄当前环境照片", data_url};
|
||||
}
|
||||
|
||||
try {
|
||||
if (!*camera)
|
||||
*camera = std::make_shared<CameraCapture>(camera_device);
|
||||
std::string jpeg = (*camera)->grab_jpeg();
|
||||
return {"已拍摄当前环境照片(" + std::to_string((*camera)->width()) + "x" +
|
||||
std::to_string((*camera)->height()) + ")",
|
||||
"data:image/jpeg;base64," + base64_encode(jpeg)};
|
||||
} catch (const std::exception& e) {
|
||||
camera->reset(); // 设备可能重新插拔,下次调用重试打开
|
||||
return {std::string("摄像头抓拍失败: ") + e.what(), ""};
|
||||
}
|
||||
};
|
||||
return t;
|
||||
}
|
||||
|
||||
Tool make_robot_status_tool() {
|
||||
Tool t;
|
||||
t.name = "get_robot_status";
|
||||
t.description = "查询机器人当前状态(电量、温度、电机、传感器)。用户询问机器人状态时调用。";
|
||||
t.parameters = json::object{{"type", "object"}, {"properties", json::object{}}};
|
||||
t.execute = [](const json::value&) -> ToolResult {
|
||||
return {"电量 85%,电机正常,温度 42°C,传感器无告警", ""};
|
||||
};
|
||||
return t;
|
||||
}
|
||||
|
||||
// ---- 命令行执行 ----
|
||||
|
||||
constexpr size_t MAX_STREAM_BYTES = 64 * 1024; // 单个流输出上限
|
||||
constexpr int MAX_CMD_TIMEOUT_S = 60;
|
||||
|
||||
struct RunResult {
|
||||
int exit_code = -1;
|
||||
bool timed_out = false;
|
||||
bool truncated = false;
|
||||
std::string out;
|
||||
std::string err;
|
||||
};
|
||||
|
||||
// 追加输出,超出上限则丢弃并标记截断
|
||||
void append_capped(std::string& dst, const char* p, size_t n, bool& truncated) {
|
||||
size_t room = MAX_STREAM_BYTES - dst.size();
|
||||
if (n > room) {
|
||||
dst.append(p, room);
|
||||
truncated = true;
|
||||
} else {
|
||||
dst.append(p, n);
|
||||
}
|
||||
}
|
||||
|
||||
RunResult run_shell(const std::string& cmd, int timeout_s) {
|
||||
int po[2], pe[2];
|
||||
if (pipe(po) != 0 || pipe(pe) != 0) throw std::runtime_error("pipe 失败");
|
||||
pid_t pid = fork();
|
||||
if (pid < 0) throw std::runtime_error("fork 失败");
|
||||
if (pid == 0) {
|
||||
dup2(po[1], STDOUT_FILENO);
|
||||
dup2(pe[1], STDERR_FILENO);
|
||||
close(po[0]);
|
||||
close(po[1]);
|
||||
close(pe[0]);
|
||||
close(pe[1]);
|
||||
execl("/bin/sh", "sh", "-c", cmd.c_str(), static_cast<char*>(nullptr));
|
||||
_exit(127);
|
||||
}
|
||||
close(po[1]);
|
||||
close(pe[1]);
|
||||
fcntl(po[0], F_SETFL, O_NONBLOCK);
|
||||
fcntl(pe[0], F_SETFL, O_NONBLOCK);
|
||||
|
||||
RunResult r;
|
||||
bool exited = false, out_eof = false, err_eof = false, killed = false;
|
||||
int status = 0;
|
||||
std::array<char, 4096> buf;
|
||||
auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(timeout_s);
|
||||
auto hard_deadline = deadline + std::chrono::seconds(3);
|
||||
|
||||
while (!(exited && out_eof && err_eof)) {
|
||||
struct pollfd fds[2] = {
|
||||
{out_eof ? -1 : po[0], POLLIN, 0},
|
||||
{err_eof ? -1 : pe[0], POLLIN, 0},
|
||||
};
|
||||
if (poll(fds, 2, 100) > 0) {
|
||||
for (const auto& fd : fds) {
|
||||
if (fd.fd < 0 || !(fd.revents & (POLLIN | POLLHUP))) continue;
|
||||
std::string& sink = (fd.fd == po[0]) ? r.out : r.err;
|
||||
ssize_t n;
|
||||
while ((n = read(fd.fd, buf.data(), buf.size())) > 0)
|
||||
append_capped(sink, buf.data(), n, r.truncated);
|
||||
if (n == 0 || (n < 0 && errno != EAGAIN && errno != EWOULDBLOCK)) {
|
||||
close(fd.fd);
|
||||
if (fd.fd == po[0]) out_eof = true;
|
||||
else err_eof = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (!exited) {
|
||||
pid_t w = waitpid(pid, &status, WNOHANG);
|
||||
if (w == pid) exited = true;
|
||||
}
|
||||
auto now = std::chrono::steady_clock::now();
|
||||
if (!killed && now >= deadline) {
|
||||
kill(pid, SIGKILL);
|
||||
killed = true;
|
||||
r.timed_out = true;
|
||||
}
|
||||
if (now >= hard_deadline) break;
|
||||
}
|
||||
|
||||
if (!out_eof) close(po[0]);
|
||||
if (!err_eof) close(pe[0]);
|
||||
if (!exited) waitpid(pid, &status, 0); // 已被 SIGKILL,等待回收
|
||||
r.exit_code = WIFEXITED(status) ? WEXITSTATUS(status) : -1;
|
||||
return r;
|
||||
}
|
||||
|
||||
std::string format_run_result(const RunResult& r) {
|
||||
std::string s = "exit_code=" + std::to_string(r.exit_code);
|
||||
if (r.timed_out) s += "(超时,已终止)";
|
||||
s += "\nstdout:\n" + (r.out.empty() ? "(空)" : r.out);
|
||||
s += "\nstderr:\n" + (r.err.empty() ? "(空)" : r.err);
|
||||
if (r.truncated) s += "\n[输出超过 64KB,已截断]";
|
||||
return s;
|
||||
}
|
||||
|
||||
Tool make_run_command_tool() {
|
||||
Tool t;
|
||||
t.name = "run_command";
|
||||
t.description =
|
||||
"在机器人主机上执行 shell 命令,返回退出码、标准输出和标准错误。"
|
||||
"可用于查看系统状态、操作硬件等。命令有超时限制,输出超过 64KB 会被截断。";
|
||||
t.parameters = json::object{
|
||||
{"type", "object"},
|
||||
{"properties", json::object{
|
||||
{"command", json::object{{"type", "string"}, {"description", "要执行的 shell 命令"}}},
|
||||
{"timeout_seconds", json::object{{"type", "integer"},
|
||||
{"description", "超时秒数,默认 10,最大 60"}}},
|
||||
}},
|
||||
{"required", json::array{"command"}},
|
||||
};
|
||||
t.execute = [](const json::value& args) -> ToolResult {
|
||||
const auto& obj = args.as_object();
|
||||
auto cit = obj.find("command");
|
||||
if (cit == obj.end() || !cit->value().is_string())
|
||||
return {"错误: 缺少 command 参数", ""};
|
||||
std::string cmd = json::value_to<std::string>(cit->value());
|
||||
|
||||
int timeout_s = 10;
|
||||
auto tit = obj.find("timeout_seconds");
|
||||
if (tit != obj.end() && tit->value().is_int64())
|
||||
timeout_s = static_cast<int>(tit->value().as_int64());
|
||||
if (timeout_s < 1) timeout_s = 1;
|
||||
if (timeout_s > MAX_CMD_TIMEOUT_S) timeout_s = MAX_CMD_TIMEOUT_S;
|
||||
|
||||
try {
|
||||
return {format_run_result(run_shell(cmd, timeout_s)), ""};
|
||||
} catch (const std::exception& e) {
|
||||
return {std::string("执行失败: ") + e.what(), ""};
|
||||
}
|
||||
};
|
||||
return t;
|
||||
}
|
||||
|
||||
// ---- 文件读写 ----
|
||||
|
||||
Tool make_read_file_tool() {
|
||||
Tool t;
|
||||
t.name = "read_file";
|
||||
t.description = "读取文本文件内容。超过 64KB 截断。";
|
||||
t.parameters = json::object{
|
||||
{"type", "object"},
|
||||
{"properties", json::object{
|
||||
{"path", json::object{{"type", "string"}, {"description", "文件路径"}}},
|
||||
}},
|
||||
{"required", json::array{"path"}},
|
||||
};
|
||||
t.execute = [](const json::value& args) -> ToolResult {
|
||||
auto it = args.as_object().find("path");
|
||||
if (it == args.as_object().end() || !it->value().is_string())
|
||||
return {"错误: 缺少 path 参数", ""};
|
||||
std::string path = json::value_to<std::string>(it->value());
|
||||
std::ifstream f(path, std::ios::binary);
|
||||
if (!f) return {"错误: 无法读取文件 " + path, ""};
|
||||
std::string content((std::istreambuf_iterator<char>(f)),
|
||||
std::istreambuf_iterator<char>());
|
||||
bool truncated = content.size() > MAX_STREAM_BYTES;
|
||||
if (truncated) content.resize(MAX_STREAM_BYTES);
|
||||
return {content + (truncated ? "\n[内容已截断]" : ""), ""};
|
||||
};
|
||||
return t;
|
||||
}
|
||||
|
||||
Tool make_write_file_tool() {
|
||||
Tool t;
|
||||
t.name = "write_file";
|
||||
t.description = "写入文本文件(覆盖已有内容)。";
|
||||
t.parameters = json::object{
|
||||
{"type", "object"},
|
||||
{"properties", json::object{
|
||||
{"path", json::object{{"type", "string"}, {"description", "文件路径"}}},
|
||||
{"content", json::object{{"type", "string"}, {"description", "要写入的内容"}}},
|
||||
}},
|
||||
{"required", json::array{"path", "content"}},
|
||||
};
|
||||
t.execute = [](const json::value& args) -> ToolResult {
|
||||
const auto& obj = args.as_object();
|
||||
auto pit = obj.find("path");
|
||||
auto cit = obj.find("content");
|
||||
if (pit == obj.end() || !pit->value().is_string() || cit == obj.end() ||
|
||||
!cit->value().is_string())
|
||||
return {"错误: 需要 path 和 content 参数", ""};
|
||||
std::string path = json::value_to<std::string>(pit->value());
|
||||
std::string content = json::value_to<std::string>(cit->value());
|
||||
std::ofstream f(path, std::ios::binary);
|
||||
if (!f) return {"错误: 无法写入文件 " + path, ""};
|
||||
f.write(content.data(), static_cast<std::streamsize>(content.size()));
|
||||
if (!f) return {"错误: 写入失败 " + path, ""};
|
||||
return {"已写入 " + std::to_string(content.size()) + " 字节到 " + path, ""};
|
||||
};
|
||||
return t;
|
||||
}
|
||||
|
||||
// ---- 子代理 ----
|
||||
|
||||
Tool make_spawn_subagent_tool(const AgentConfig& cfg) {
|
||||
Tool t;
|
||||
t.name = "spawn_subagent";
|
||||
t.description =
|
||||
"委派任务给子代理执行。子代理拥有扩展技能:执行 shell 命令、读写文件、拍照、"
|
||||
"查询状态。当任务需要运行命令、读写文件等主代理不具备的能力时调用。";
|
||||
t.parameters = json::object{
|
||||
{"type", "object"},
|
||||
{"properties", json::object{
|
||||
{"task", json::object{{"type", "string"}, {"description", "委派给子代理的具体任务"}}},
|
||||
}},
|
||||
{"required", json::array{"task"}},
|
||||
};
|
||||
t.execute = [cfg](const json::value& args) -> ToolResult {
|
||||
auto it = args.as_object().find("task");
|
||||
if (it == args.as_object().end() || !it->value().is_string())
|
||||
return {"错误: 缺少 task 参数", ""};
|
||||
std::string task = json::value_to<std::string>(it->value());
|
||||
try {
|
||||
return {run_subagent(cfg, task), ""};
|
||||
} catch (const std::exception& e) {
|
||||
return {std::string("子代理执行失败: ") + e.what(), ""};
|
||||
}
|
||||
};
|
||||
return t;
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
#pragma once
|
||||
|
||||
#include <boost/json.hpp>
|
||||
|
||||
#include <functional>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace json = boost::json;
|
||||
|
||||
struct AgentConfig; // 前向声明(定义见 agent.hpp),避免头文件循环依赖
|
||||
|
||||
// 工具执行结果。image_data_url 非空表示工具返回了一张图片,
|
||||
// 由 agent 主循环把图片注入下一轮对话(OpenAI 协议的 tool 消息只允许文本)。
|
||||
struct ToolResult {
|
||||
std::string text;
|
||||
std::string image_data_url;
|
||||
};
|
||||
|
||||
struct Tool {
|
||||
std::string name;
|
||||
std::string description;
|
||||
json::value parameters; // JSON Schema(object)
|
||||
std::function<ToolResult(const json::value& args)> execute;
|
||||
};
|
||||
|
||||
class ToolRegistry {
|
||||
public:
|
||||
void add(Tool tool) { tools_.push_back(std::move(tool)); }
|
||||
|
||||
const std::vector<Tool>& all() const { return tools_; }
|
||||
|
||||
bool execute(const std::string& name, const json::value& args, ToolResult& out) const {
|
||||
for (const auto& t : tools_) {
|
||||
if (t.name == name) {
|
||||
out = t.execute(args);
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
private:
|
||||
std::vector<Tool> tools_;
|
||||
};
|
||||
|
||||
// 拍照工具:simulated_image_path 非空时读取该文件(模拟模式),
|
||||
// 否则从真实摄像头设备(V4L2,默认 /dev/video0)抓帧。
|
||||
Tool make_capture_image_tool(const std::string& simulated_image_path,
|
||||
const std::string& camera_device);
|
||||
|
||||
// 状态查询工具:返回模拟的机器人状态。
|
||||
Tool make_robot_status_tool();
|
||||
|
||||
// 命令行工具:执行 shell 命令(带超时与输出截断)。
|
||||
Tool make_run_command_tool();
|
||||
|
||||
// 文件工具:读取/写入文本文件。
|
||||
Tool make_read_file_tool();
|
||||
Tool make_write_file_tool();
|
||||
|
||||
// 子代理工具:把任务委派给拥有扩展技能的子代理执行。
|
||||
Tool make_spawn_subagent_tool(const AgentConfig& cfg);
|
||||
+1
-1
@@ -144,7 +144,7 @@ waf/
|
||||
│ ├── engine/
|
||||
│ │ └── engine.go # coraza 封装:加载/热加载/列举/单测/输入检测
|
||||
│ ├── logclient/
|
||||
│ │ └── logclient.go # Go 版 pymsg 日志客户端(异步批量推送)
|
||||
│ │ └── logclient.go # gomsg 日志客户端适配层(协议/队列实现见 ../gomsg)
|
||||
│ └── proxy/
|
||||
│ └── proxy.go # ReverseProxy + WAF 事务接入 + 安全响应头注入
|
||||
├── rules/
|
||||
|
||||
@@ -4,9 +4,12 @@ go 1.22.0
|
||||
|
||||
require (
|
||||
github.com/corazawaf/coraza/v3 v3.3.0
|
||||
github.com/micro_servers/gomsg v0.0.0
|
||||
gopkg.in/yaml.v3 v3.0.1
|
||||
)
|
||||
|
||||
replace github.com/micro_servers/gomsg => ../gomsg
|
||||
|
||||
require (
|
||||
github.com/corazawaf/libinjection-go v0.2.2 // indirect
|
||||
github.com/magefile/mage v1.15.1-0.20241126214340-bdc92f694516 // indirect
|
||||
|
||||
@@ -1,142 +1,38 @@
|
||||
// Package logclient 提供异步批量日志推送客户端。
|
||||
// Package logclient 基于 gomsg 的集中式日志客户端(异步批量推送 db2 Logs 表)。
|
||||
//
|
||||
// 通过 pymsg TCP 协议调用 dbv2 execute 路由,逐条 INSERT 到 Logs 表。
|
||||
//
|
||||
// 设计要点(与 Python example_log_client.py 对齐):
|
||||
// - 异步非阻塞:内存 channel + 后台 goroutine 批量推送,不拖慢业务
|
||||
// - fail-safe:db 服务不可用时静默丢弃,绝不影响业务
|
||||
// - 缓冲批量:队列满 50 条或超时 2 秒时自动 flush
|
||||
//
|
||||
// 协议:pymsg 长度前缀帧(4 字节 big-endian uint32 长度 + UTF-8 载荷)
|
||||
// 原实现(pymsg 帧协议 + 异步队列)已抽取至 ../gomsg,本包保留包级 API
|
||||
// 兼容层:main / engine / proxy 的调用点零改动。
|
||||
package logclient
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"os"
|
||||
"sync"
|
||||
"time"
|
||||
import "github.com/micro_servers/gomsg"
|
||||
|
||||
var (
|
||||
lc *gomsg.LogClient
|
||||
pendingLevel string
|
||||
)
|
||||
|
||||
const (
|
||||
defaultDBAddr = "127.0.0.1:8005"
|
||||
batchSize = 50
|
||||
flushInterval = 2 * time.Second
|
||||
socketTimeout = 5 * time.Second
|
||||
// 消息大小安全上限(64 MiB),与 Python / C++ 端保持一致。
|
||||
maxMsgSize = 64 * 1024 * 1024
|
||||
)
|
||||
|
||||
// 日志级别数值(越小越详细)。
|
||||
const (
|
||||
levelDebug = 0
|
||||
levelInfo = 1
|
||||
levelWarn = 2
|
||||
levelError = 3
|
||||
)
|
||||
|
||||
var levelMap = map[string]int{
|
||||
"debug": levelDebug,
|
||||
"info": levelInfo,
|
||||
"warn": levelWarn,
|
||||
"warning": levelWarn,
|
||||
"error": levelError,
|
||||
"DEBUG": levelDebug,
|
||||
"INFO": levelInfo,
|
||||
"WARN": levelWarn,
|
||||
"WARNING": levelWarn,
|
||||
"ERROR": levelError,
|
||||
}
|
||||
|
||||
// dbAddr 通过环境变量 LOG_DB_ADDR 配置(默认 127.0.0.1:8005),支持非标准部署场景。
|
||||
var dbAddr = defaultDBAddr
|
||||
|
||||
// minLevel 控制最低输出级别,默认 info(debug 被过滤)。
|
||||
var minLevel = levelInfo
|
||||
|
||||
func init() {
|
||||
if v := os.Getenv("LOG_DB_ADDR"); v != "" {
|
||||
dbAddr = v
|
||||
}
|
||||
if v := os.Getenv("LOG_LEVEL"); v != "" {
|
||||
SetLevel(v)
|
||||
// Init 初始化全局日志客户端(读 LOG_DB_ADDR / LOG_LEVEL 环境变量)。
|
||||
func Init() {
|
||||
lc = gomsg.NewLogClient("")
|
||||
if pendingLevel != "" {
|
||||
lc.SetLevel(pendingLevel)
|
||||
}
|
||||
}
|
||||
|
||||
// SetLevel 设置最低日志级别(debug|info|warn|error,大小写不敏感)。
|
||||
func SetLevel(level string) {
|
||||
if n, ok := levelMap[level]; ok {
|
||||
minLevel = n
|
||||
pendingLevel = level
|
||||
if lc != nil {
|
||||
lc.SetLevel(level)
|
||||
}
|
||||
}
|
||||
|
||||
// Entry 单条日志。
|
||||
type Entry struct {
|
||||
Service string `json:"service"`
|
||||
Level string `json:"level"`
|
||||
Message string `json:"message"`
|
||||
TraceID string `json:"trace_id,omitempty"`
|
||||
Extra map[string]any `json:"extra,omitempty"`
|
||||
}
|
||||
|
||||
// dbV2Request execute 路由的 INSERT 请求体。
|
||||
type dbV2Request struct {
|
||||
Route string `json:"_route"`
|
||||
Action string `json:"action"`
|
||||
Table string `json:"table"`
|
||||
Values map[string]any `json:"values"`
|
||||
}
|
||||
|
||||
// dbV2Response execute 路由的通用响应。
|
||||
type dbV2Response struct {
|
||||
Success bool `json:"success"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// Client 异步批量日志推送客户端。
|
||||
// send 仅在 flushLoop 单 goroutine 中调用,conn 复用无需额外加锁。
|
||||
type Client struct {
|
||||
ch chan Entry
|
||||
stop chan struct{}
|
||||
wg sync.WaitGroup
|
||||
stopped bool
|
||||
mu sync.Mutex
|
||||
conn net.Conn // 复用的 pymsg 长连接,失败时重建
|
||||
}
|
||||
|
||||
// global 全局单例。
|
||||
var global *Client
|
||||
|
||||
// Init 初始化全局日志客户端并启动后台 flush goroutine。
|
||||
func Init() {
|
||||
global = New()
|
||||
}
|
||||
|
||||
// New 创建客户端实例。
|
||||
func New() *Client {
|
||||
c := &Client{
|
||||
ch: make(chan Entry, 256),
|
||||
stop: make(chan struct{}),
|
||||
}
|
||||
c.wg.Add(1)
|
||||
go c.flushLoop()
|
||||
return c
|
||||
}
|
||||
|
||||
// Push 推送一条日志(非阻塞,队列满时丢弃)。
|
||||
// 低于 SetLevel 设定的最低级别的日志将被静默丢弃。
|
||||
// Push 推送一条日志(非阻塞,队列满或低于最低级别时静默丢弃)。
|
||||
func Push(service, level, message string, extra map[string]any) {
|
||||
if global == nil {
|
||||
if lc == nil {
|
||||
return
|
||||
}
|
||||
n, ok := levelMap[level]
|
||||
if !ok || n < minLevel {
|
||||
return
|
||||
}
|
||||
global.push(Entry{
|
||||
lc.Push(gomsg.Entry{
|
||||
Service: service,
|
||||
Level: level,
|
||||
Message: message,
|
||||
@@ -164,199 +60,9 @@ func Error(service, message string, extra map[string]any) {
|
||||
Push(service, "ERROR", message, extra)
|
||||
}
|
||||
|
||||
// Flush 同步 flush 队列中剩余日志(用于进程退出前)。
|
||||
// Flush 停止队列并 flush 剩余日志(用于进程退出前)。
|
||||
func Flush() {
|
||||
if global == nil {
|
||||
return
|
||||
}
|
||||
global.flush()
|
||||
}
|
||||
|
||||
func (c *Client) push(e Entry) {
|
||||
select {
|
||||
case c.ch <- e:
|
||||
default:
|
||||
// 队列满,静默丢弃(fail-safe)
|
||||
if lc != nil {
|
||||
lc.Flush()
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) flushLoop() {
|
||||
defer c.wg.Done()
|
||||
batch := make([]Entry, 0, batchSize)
|
||||
ticker := time.NewTicker(flushInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case e := <-c.ch:
|
||||
batch = append(batch, e)
|
||||
// 非阻塞地继续取,直到满 batchSize
|
||||
for len(batch) < batchSize {
|
||||
select {
|
||||
case e2 := <-c.ch:
|
||||
batch = append(batch, e2)
|
||||
default:
|
||||
goto send
|
||||
}
|
||||
}
|
||||
send:
|
||||
if len(batch) >= batchSize {
|
||||
c.send(batch)
|
||||
batch = batch[:0]
|
||||
}
|
||||
case <-ticker.C:
|
||||
if len(batch) > 0 {
|
||||
c.send(batch)
|
||||
batch = batch[:0]
|
||||
}
|
||||
case <-c.stop:
|
||||
// 退出前 drain channel
|
||||
for {
|
||||
select {
|
||||
case e := <-c.ch:
|
||||
batch = append(batch, e)
|
||||
default:
|
||||
if len(batch) > 0 {
|
||||
c.send(batch)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) send(batch []Entry) {
|
||||
now := time.Now().UTC().Format(time.RFC3339)
|
||||
for _, entry := range batch {
|
||||
c.sendOne(entry, now)
|
||||
}
|
||||
}
|
||||
|
||||
// sendOne 通过 pymsg TCP 协议发送单条日志到 dbv2。
|
||||
func (c *Client) sendOne(entry Entry, ts string) {
|
||||
extraJSON := ""
|
||||
if entry.Extra != nil {
|
||||
b, err := json.Marshal(entry.Extra)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
extraJSON = string(b)
|
||||
}
|
||||
|
||||
values := map[string]any{
|
||||
"ts": ts,
|
||||
"service": entry.Service,
|
||||
"level": entry.Level,
|
||||
"message": entry.Message,
|
||||
}
|
||||
if entry.TraceID != "" {
|
||||
values["trace_id"] = entry.TraceID
|
||||
}
|
||||
if extraJSON != "" {
|
||||
values["extra"] = extraJSON
|
||||
}
|
||||
|
||||
req := dbV2Request{
|
||||
Route: "execute",
|
||||
Action: "insert",
|
||||
Table: "Logs",
|
||||
Values: values,
|
||||
}
|
||||
data, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// pymsg 协议:4 字节 big-endian 长度 + UTF-8 载荷(复用长连接,失败时重连一次)
|
||||
resp, err := c.pymsgCall(data)
|
||||
if err != nil {
|
||||
return // fail-safe
|
||||
}
|
||||
|
||||
var r dbV2Response
|
||||
if json.Unmarshal(resp, &r) != nil || !r.Success {
|
||||
fmt.Fprintf(os.Stderr, "[logclient] INSERT failed: %s\n", r.Error)
|
||||
}
|
||||
}
|
||||
|
||||
// pymsgCall 复用长连接发送一条 pymsg 请求并读取响应。
|
||||
// 首次失败时关闭旧连接重建一次,避免高并发拦截场景下大量 TIME_WAIT 短连接。
|
||||
func (c *Client) pymsgCall(payload []byte) ([]byte, error) {
|
||||
resp, err := c.doPymsg(payload)
|
||||
if err == nil {
|
||||
return resp, nil
|
||||
}
|
||||
// 旧连接可能已断开(db 重启等),关闭后重连一次
|
||||
c.closeConn()
|
||||
return c.doPymsg(payload)
|
||||
}
|
||||
|
||||
// doPymsg 在当前连接上执行一次请求-响应交互。
|
||||
func (c *Client) doPymsg(payload []byte) ([]byte, error) {
|
||||
conn, err := c.getConn()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
conn.SetDeadline(time.Now().Add(socketTimeout))
|
||||
|
||||
// 发送:4 字节长度前缀 + 载荷
|
||||
header := make([]byte, 4)
|
||||
binary.BigEndian.PutUint32(header, uint32(len(payload)))
|
||||
if _, err := conn.Write(header); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if _, err := conn.Write(payload); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 接收:4 字节长度前缀 + 载荷
|
||||
if _, err := io.ReadFull(conn, header); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
respLen := binary.BigEndian.Uint32(header)
|
||||
if respLen > maxMsgSize {
|
||||
return nil, fmt.Errorf("response too large: %d bytes (max %d)", respLen, maxMsgSize)
|
||||
}
|
||||
resp := make([]byte, respLen)
|
||||
if _, err := io.ReadFull(conn, resp); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
// getConn 返回当前复用连接,不存在时建立。
|
||||
func (c *Client) getConn() (net.Conn, error) {
|
||||
if c.conn != nil {
|
||||
return c.conn, nil
|
||||
}
|
||||
conn, err := net.DialTimeout("tcp", dbAddr, socketTimeout)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.conn = conn
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
// closeConn 关闭并置空复用连接。
|
||||
func (c *Client) closeConn() {
|
||||
if c.conn != nil {
|
||||
_ = c.conn.Close()
|
||||
c.conn = nil
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) flush() {
|
||||
c.mu.Lock()
|
||||
if c.stopped {
|
||||
c.mu.Unlock()
|
||||
return
|
||||
}
|
||||
c.stopped = true
|
||||
c.mu.Unlock()
|
||||
|
||||
close(c.stop)
|
||||
c.wg.Wait()
|
||||
c.closeConn()
|
||||
}
|
||||
+76
-2
@@ -12,12 +12,15 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"encoding/json"
|
||||
"flag"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/signal"
|
||||
@@ -200,9 +203,37 @@ func main() {
|
||||
logclient.Info("waf", fmt.Sprintf("服务启动 instance=%s listen=%s upstream=%s rules=%s pid=%d",
|
||||
*instance, cfg.Listen, cfg.Upstream, cfg.RulesDir, os.Getpid()), nil)
|
||||
|
||||
errCh := make(chan error, 1)
|
||||
errCh := make(chan error, 2)
|
||||
// 同一端口区分 TLS 与明文 HTTP:明文请求 301 重定向到 HTTPS(原 redirect80 的职责)
|
||||
cert, err := tls.LoadX509KeyPair(cfg.TLSCert, cfg.TLSKey)
|
||||
if err != nil {
|
||||
logclient.Error("waf", fmt.Sprintf("加载 TLS 证书失败: %v", err), nil)
|
||||
logclient.Flush()
|
||||
os.Exit(1)
|
||||
}
|
||||
server.TLSConfig.Certificates = []tls.Certificate{cert}
|
||||
rawLn, err := net.Listen("tcp", cfg.Listen)
|
||||
if err != nil {
|
||||
logclient.Error("waf", fmt.Sprintf("监听失败: %v", err), nil)
|
||||
logclient.Flush()
|
||||
os.Exit(1)
|
||||
}
|
||||
tlsCh := make(chan net.Conn)
|
||||
plainCh := make(chan net.Conn)
|
||||
go acceptAndDispatch(rawLn, tlsCh, plainCh)
|
||||
|
||||
plainServer := &http.Server{
|
||||
Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
http.Redirect(w, r, "https://"+r.Host+r.URL.RequestURI(), http.StatusMovedPermanently)
|
||||
}),
|
||||
ReadHeaderTimeout: 10 * time.Second,
|
||||
}
|
||||
|
||||
go func() {
|
||||
errCh <- server.ListenAndServeTLS(cfg.TLSCert, cfg.TLSKey)
|
||||
errCh <- server.Serve(tls.NewListener(&chanListener{ch: tlsCh, addr: rawLn.Addr()}, server.TLSConfig))
|
||||
}()
|
||||
go func() {
|
||||
errCh <- plainServer.Serve(&chanListener{ch: plainCh, addr: rawLn.Addr()})
|
||||
}()
|
||||
|
||||
select {
|
||||
@@ -219,10 +250,53 @@ func main() {
|
||||
if err := server.Shutdown(ctx); err != nil {
|
||||
logclient.Error("waf", fmt.Sprintf("优雅关闭超时: %v", err), nil)
|
||||
}
|
||||
plainServer.Shutdown(ctx)
|
||||
logclient.Flush()
|
||||
}
|
||||
}
|
||||
|
||||
// bufferedConn 保留已 peek 的首字节,避免后续 Read 丢数据。
|
||||
type bufferedConn struct {
|
||||
net.Conn
|
||||
r *bufio.Reader
|
||||
}
|
||||
|
||||
func (c *bufferedConn) Read(p []byte) (int, error) { return c.r.Read(p) }
|
||||
|
||||
// chanListener 从 channel 取连接,供两个 http.Server 分别消费 TLS / 明文连接。
|
||||
type chanListener struct {
|
||||
ch chan net.Conn
|
||||
addr net.Addr
|
||||
}
|
||||
|
||||
func (l *chanListener) Accept() (net.Conn, error) {
|
||||
c, ok := <-l.ch
|
||||
if !ok {
|
||||
return nil, io.EOF
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (l *chanListener) Close() error { return nil }
|
||||
func (l *chanListener) Addr() net.Addr { return l.addr }
|
||||
|
||||
// acceptAndDispatch 按首字节分发:0x16 为 TLS ClientHello → TLS 通道,其余 → 明文 HTTP 通道。
|
||||
func acceptAndDispatch(ln net.Listener, tlsCh, plainCh chan<- net.Conn) {
|
||||
for {
|
||||
c, err := ln.Accept()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
br := bufio.NewReader(c)
|
||||
b, err := br.Peek(1)
|
||||
if err == nil && len(b) == 1 && b[0] == 0x16 {
|
||||
tlsCh <- &bufferedConn{Conn: c, r: br}
|
||||
} else {
|
||||
plainCh <- &bufferedConn{Conn: c, r: br}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func loadConfig(path string) *Config {
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
|
||||
Reference in new issue
Block a user