12 Commits

Author SHA1 Message Date
wunianxiao 3476d8cbcb Added tool implementation for reading tools. 2026-06-23 20:36:16 +08:00
wunianxiao 132089c8eb feat(context): /clear 指令截断上下文并按会话过期归档
用户发送 /clear 时主动截断上下文:加载历史 → LLM 生成摘要归档
→ 丢弃在途会话 → 回复确认(带摘要预览)。

- context/types: SummaryReason 新增 Clear 变体
- context/builder: 抽出 archive_session 共享逻辑,
  新增 trigger_clear_summary(返回摘要文本)
- db/models: 新增 latest_session_boundary;list_recent_chat_records
  _for_context 按会话边界过滤(clear/timeout),并排除 source=clear
- daemon: SessionMap 记录 user_id;拦截 /clear 走 handle_clear_command;
  新增 build_clear_reply 纯函数
- migration: session_summaries (user_id, created_at DESC) 索引

会话边界机制统一 clear 与 12h 空闲超时;摘要不注入下一对话,
只存档供 read_summaries 查询。新增 4 个单元测试。
2026-06-23 11:57:07 +08:00
wunianxiao 0299329fb2 Merge branch 'security/sandbox-boundaries'
两层安全边界落地:
- 授权层:统一工具级审批,堵住 stdio→result 绕过
- 隔离层:长期沙箱容器 + docker exec,仅不受信任代码进容器
- 审计层:tool_manager 调用全留痕

9 个提交,28 测试通过,clippy 无新增警告。
2026-06-22 17:51:58 +08:00
wunianxiao a2f88b9b01 fix(sandbox): 容器不可用自动降级 + clean_expired 并发回归测试
P1-1 降级机制:
  - sandbox 新增 DEGRADED 全局标志,ensure_container 失败时置位
  - 新增 ensure_ready_or_degrade: exec 前检测容器运行状态,
    不在则重建,重建失败置降级标志并返回 false
  - executor 调用 ensure_ready_or_degrade,降级时回退本地执行
    (审批已在入口完成,安全边界不受影响)
  - spawn_child 保留 spawn 级失败兜底(docker 二进制缺失等)
  - 解决: IAS_DOCKER_SANDBOX=1 但 docker 异常时工具不再持续 SpawnFailed

P1-2 并发回归测试:
  - approval 新增 create_with_ttl 测试入口(cfg(test))
  - clean_expired_only_removes_expired_entries: 验证只删过期条目
  - clean_expired_safe_when_no_pending: 空 pending 不 panic
  - clean_expired_skips_entries_already_removed_by_reply:
    回归保护——clean_expired 第一阶段收集 hash 后、第二阶段删除前,
    条目被 handle_reply 提前移除时,按 hash 查找应安全跳过不误删

P2-5 容器崩溃恢复:
  - ensure_ready_or_degrade 统一处理容器崩溃重建
  - recover_if_crashed 保留供 spawn 级失败兜底调用

P2-6 build_exec_cmd 用 spec.tool 替代 path 文件名提取,消除隐式耦合

验证: 28 测试通过(原25+3),clippy 22<基线23 无新增警告,
端到端验证容器重建/降级/本地三种路径均正确。
2026-06-22 17:46:48 +08:00
wunianxiao 142efb070f refactor(sandbox): 长期容器 + docker exec 替代一次性 docker run
容器策略从"每次 docker run 一次性容器"改为"长期容器 + docker exec":

- 统一产物目录 tools/bin/:build.sh 和 tool_manager 编译后复制产物到
  此目录,registry 解析路径优先用 tools/bin/<tool>。
- 长期容器 ias-sandbox:daemon 启动时 ensure(已运行则复用),
  tools/bin 只读挂载到 /tools,无网络/只读根/cap-drop ALL/资源限制。
- docker exec 调用:工具执行改为 docker exec -i ias-sandbox
  timeout -k 5 <secs> /tools/<name>,省去容器创建开销。
  timeout -k 5 确保超时 SIGTERM→SIGKILL,无孤儿进程。
- 新增 src/tools/sandbox.rs:ensure_container + build_exec_cmd。
- executor 删除 build_docker_cmd/sandbox_enabled,docker 分支改用
  sandbox::build_exec_cmd。
- env 按 spec 白名单 -e 注入(容器启动时仅有 PATH)。

验证:本地执行 hostname=debian,docker exec hostname=容器ID,
证明确实走长期容器 exec;25 测试通过。
2026-06-22 17:11:05 +08:00
wunianxiao 092f3b78f2 refactor(sandbox): 按威胁模型精确化——仅不受信任代码进容器
http/db 类工具的子进程是瘦壳(纯计算),危险操作全在受信任主进程,
套容器是纯开销无隔离收益。调整为:

- 开发者预写工具:移除 spec 的 sandbox: profile: stdio,回默认 local
  本地执行(仍保留 env_clear + 白名单注入)
- tool_manager 产出的工具:spec_yaml 自动写 sandbox: profile: stdio,
  不受信任代码(LLM 提供 rust_code)默认进 stdio 容器
- tool_manager 编译:仍走 builder 容器
- 移除 ias-network-runner Dockerfile(当前无工具直接联网)
- executor docker 分支仅保留 stdio/builder,删 network 分支

验证:datetime 本地执行 +08:00;tool_manager 产出工具 spec 含
sandbox: profile: stdio 且容器内执行成功;25 测试通过。
2026-06-22 16:53:36 +08:00
wunianxiao 4a5b21a69e docs: 两层安全边界设计文档
授权边界(统一审批)+ 隔离边界(Docker 一次性容器)+ 审计
的设计、实现与验证记录。
2026-06-22 11:51:05 +08:00
wunianxiao cacd4436ce fix(sandbox): builder 容器 --user uid:gid + 离线模式,修 cap-drop 写权限
步骤 9: tool_manager build 迁移到 builder 容器验证通过。

问题: --cap-drop ALL 丢掉 CAP_DAC_OVERRIDE,容器内 root 无法写
      bind mount 的非 root 目录(宿主 xiao:1000)。
      这是 cap-drop 正确生效的安全行为,非 bug。

修复:
- builder 容器 --user uid:gid 以宿主用户身份运行,自然能写自己
  的目录,同时保留 cap-drop ALL + no-new-privileges。
- Cargo.toml 改 CARGO_HOME=/tmp/cargo + chmod 777,使任意 uid
  都能复用 serde_json 预缓存。
- cargo build --offline + CARGO_NET_OFFLINE=true,配合 --network
  none 真正离线编译,杜绝恶意代码下载依赖。

端到端: tool_manager create_tool → builder 容器离线编译 →
产物运行成功,审计日志记录正确。
2026-06-22 11:49:31 +08:00
wunianxiao 20f7c3ba9c feat(sandbox): Docker 一次性容器执行隔离边界(步骤 6-7)
引入两类安全边界:统一审批=授权边界(上一提交),一次性 Docker
容器=执行隔离边界(本提交)。LLM 即使执行恶意代码也只能在受限容器
内运行。

spec.rs: 新增 SandboxSpec(profile/network/memory/cpus/pids_limit/
  read_only),默认 profile=local 不启用容器。
executor.rs: 按 sandbox.profile 选择执行后端
  - local: 本地 spawn(env_clear + 白名单注入)
  - stdio: ias-stdio-runner,--network none --read-only --cap-drop ALL
    --security-opt no-new-privileges,仅 bind mount 工具二进制(ro)+tmpfs /tmp
  - network: ias-network-runner(预留自行联网工具)
  - builder: ias-tool-builder(供 tool_manager 编译)
  全局开关 IAS_DOCKER_SANDBOX=1,未启用时行为不变(向后兼容)。
tool_manager: build_tool 支持 IAS_DOCKER_BUILDER=1 在 ias-tool-builder
  容器内离线编译(--network none,挂载 tool_dir 到 /work)。
docker/: 三类 Dockerfile + build-images.sh 构建脚本。
  - stdio-runner: debian-slim + ca-certificates,无网络只读
  - network-runner: 预留联网工具
  - tool-builder: rust:1.88-slim 预缓存 serde_json,支持离线编译
spec 迁移: 17 个普通工具追加 sandbox.profile: stdio;tool_manager
  保持 local(自身在宿主操作文件系统)。

端到端验证: datetime 工具本地模式 +08:00,Docker 模式 +00:00(UTC),
  --rm 无容器残留。
2026-06-22 11:38:42 +08:00
wunianxiao 1e33f527e0 fix(approval): clean_expired 按 code_hash 定位删除,修复并发 index 错位
P0-5: 原实现快照 (uid, index, code_hash) 后释放锁更新 DB,再重新加锁
      按 index remove。若两阶段间 handle_reply 增删了条目,index 指向
      错误条目,可能把未过期的审批误标 Expired 并通知等待方。

      改为两阶段都按 code_hash 精确定位:收集过期 hash → 更新 DB →
      重新加锁按 hash 查找删除,并发回复只会让查找安全跳过。
2026-06-22 10:48:16 +08:00
wunianxiao 4f80e81b0b security(executor): 统一工具级审批 + env 白名单 + 结构化结果
P0-1 统一审批: 高风险工具在 spawn 前统一走审批,与后续消息类型
      (result/http/db) 无关,堵住 stdio→result 绕过审批的漏洞。
      移除 http/db 分支内重复审批(入口已授权)。
P0-4 env 白名单: spawn 前 cmd.env_clear(),仅注入 PATH + spec 声明
      的 env 变量,避免工具读取宿主全部 API key/数据库密码。
      tool_manager spec 补充 HOME/CARGO_HOME/RUSTUP_HOME 供 cargo 使用。
P1-10 结构化结果: 引入 ToolRunOutcome{Ok,Err{kind,message}} 与
      ToolErrorKind 枚举,executor/registry dispatch 返回结构化结果。
      daemon 与 CLI 据此判断成败,删除脆弱的 is_tool_error 字符串前缀
      匹配;失败时记录 kind 便于追踪。
2026-06-22 10:47:01 +08:00
wunianxiao bf231352fb security(tool_manager): 禁止自改保留名 + 强制 high risk + 审计日志 + 构建超时清理
P0-2: validate_tool_name 将 tool_manager 等列为保留名,禁止
      create/update/build 作用于自身,杜绝持久化后门。
P0-3: spec_yaml 忽略调用方传入的 risk_level,由 tool_manager
      产出的工具一律 high risk,避免创建低风险后门绕过审批。
审计: 每次调用追加记录到 tools/.tool_audit.log
      (timestamp/user_id/action/tool_name/ok|err/detail)。
构建: 用 timeout 限制 180s,失败时删除残留旧二进制,
      避免 reload 后调用旧版本造成"假成功"。
2026-06-22 10:37:42 +08:00
43 changed files with 1786 additions and 211 deletions
+3
View File
@@ -1,11 +1,14 @@
# Rust 构建产物 # Rust 构建产物
/target /target
tools/*/target/ tools/*/target/
tools/bin/
dist/ dist/
# 运行时数据(状态文件、缓存、聊天记录本地副本等) # 运行时数据(状态文件、缓存、聊天记录本地副本等)
.data/ .data/
.pi .pi
# tool_manager 审计日志
tools/.tool_audit.log
# 环境配置(含密钥,不入库) # 环境配置(含密钥,不入库)
.env .env
+27
View File
@@ -333,6 +333,33 @@ flowchart TB
I --> J["发送给 LLM"] I --> J["发送给 LLM"]
``` ```
### 5.3.1 `/clear` 清空上下文流程
用户发送 `/clear` 指令时,主动截断上下文并把本次会话归档为摘要(按会话过期处理):
```mermaid
flowchart TB
U(["📱 用户发送 /clear"]) --> K{"is_clear_command() 匹配?"}
K -->|"否"| N(["走正常 LLM 对话流程"])
K -->|"是"| L["加载历史(过滤掉 /clear 本身)"]
L --> M{"历史非空?"}
M -->|"否"| R2["回复: 上下文已是空的"]
M -->|"是"| S["一次性 Conversation + summarizer"]
S --> T["trigger_clear_summary() → LLM 生成摘要"]
T --> DB[("session_summaries\nreason=clear")]
DB --> D["丢弃该用户在途会话 (sessions.retain)"]
D --> R1["回复确认 + 摘要预览"]
R2 --> END(["📱 用户收到回复"])
R1 --> END
```
- **会话边界**:摘要的 `created_at` 作为边界,之后加载历史时通过
`latest_session_boundary()` 过滤掉此前的消息(`/clear` 与 12h 空闲超时
共用同一套边界机制)
- **注入与否**`/clear` 摘要**不会**注入下一次对话的 system prompt
只存档供 `read_summaries` 工具查询(与空闲超时摘要一致)
- **在途会话**`/clear` 会丢弃该用户尚未完成的工具往返,避免用旧上下文继续
--- ---
## 6. 文件依赖矩阵 ## 6. 文件依赖矩阵
+7
View File
@@ -0,0 +1,7 @@
# ias-stdio-runner — 运行普通 stdio 工具的极简隔离环境
# 无网络、只读根、低权限。工具二进制通过 bind mount 提供,镜像本身不含工具。
FROM debian:bookworm-slim
RUN apt-get update \
&& apt-get install -y --no-install-recommends ca-certificates \
&& rm -rf /var/lib/apt/lists/*
ENTRYPOINT []
+13
View File
@@ -0,0 +1,13 @@
# ias-tool-builder — tool_manager 编译 Rust 工具的隔离环境
# 预缓存 serde_jsontool_manager 生成工具的唯一依赖),支持 --network none 离线编译。
# CARGO_HOME=/tmp/cargo 并 chmod 777,使任意 uid--user 运行)都能复用缓存。
# 运行时由 tool_manager 挂载 tool_dir 到 /workcargo 在 /work 内编译,产物写回宿主。
FROM rust:1.88-slim
ENV CARGO_HOME=/tmp/cargo
RUN cargo new /tmp/seed && cd /tmp/seed \
&& printf 'serde_json = "1.0"\n' >> Cargo.toml \
&& cargo build --release \
&& rm -rf /tmp/seed \
&& chmod -R 777 /tmp/cargo
WORKDIR /work
ENTRYPOINT []
+16
View File
@@ -0,0 +1,16 @@
#!/usr/bin/env bash
# 构建沙箱基础镜像。
# - ias-stdio-runner : 运行 tool_manager 产出的不受信任工具(无网络只读)
# - ias-tool-builder : tool_manager 编译 Rust 工具(离线)
# 无需为每个工具单独写 Dockerfile。
set -euo pipefail
cd "$(dirname "$0")" # docker/
echo "=== 构建 ias-stdio-runner(无网络只读容器)==="
docker build -t ias-stdio-runner -f Dockerfile.stdio-runner .
echo "=== 构建 ias-tool-builderRust 编译容器)==="
docker build -t ias-tool-builder -f Dockerfile.tool-builder .
echo "=== 完成 ==="
docker images | grep -E '^ias-(stdio|tool)'
+186
View File
@@ -0,0 +1,186 @@
# 两层安全边界设计
> 审批决定"能不能做"Docker 限制"做坏了能坏到哪里",审计回答"谁让系统新增或修改了什么能力"。
## 架构
```
LLM 调用工具
→ registry.dispatch(name, params, user_id, approved)
→ load spec + validate params
→ 【授权边界】if high_risk && !approved: 统一审批(不依赖消息类型)
→ 【隔离边界】按 sandbox.profile 选择执行方式:
- local : 本地 spawnenv_clear + 白名单注入)
- stdio : docker run ias-stdio-runner (无网络、只读根)
- network : docker run ias-network-runner(联网、资源限制)
- builder : docker run ias-tool-builder (离线编译,仅 tool_manager
→ stdin 传 envelope → stdout 读 JSON → 超时 kill → 容器销毁
```
## 第一层:统一工具级审批(授权边界)
**位置**`src/tools/executor.rs` `execute_loop()` 入口
高风险工具在 **spawn 之前** 必须获得用户确认,与工具后续输出的消息类型(`result`/`http`/`db`)无关。
```rust
if spec.is_high_risk() && !approved {
// 发送确认码 → 等待用户回复 → Approved/Rejected/Expired
}
```
**修复的漏洞**:原实现审批只在 `http`/`db` 消息分支触发,`stdio → result` 可绕过。`tool_manager`high risk)直接输出 result 即可零确认执行。
## 第二层:长期沙箱容器 + docker exec(隔离边界)
### 适用范围(按威胁模型精确化)
http/db 类工具的子进程是"瘦壳"——危险操作(联网、DB 读写)全在受信任的主进程里,
给它们套容器是纯开销,**没有隔离收益**。因此:
| 场景 | 是否容器隔离 | 原因 |
|------|------------|------|
| 开发者预写工具(weather/amap/... | ❌ 本地执行 | 子进程纯计算,危险操作在主进程 |
| tool_manager 产出的工具 | ✅ stdio 容器 | 含 LLM 提供的 rust_code,不受信任 |
| tool_manager 编译 | ✅ builder 容器 | cargo 执行任意 build script/proc macro |
### 统一产物目录
所有工具二进制统一收集到 `tools/bin/``build.sh` 编译后复制,`tool_manager` 编译后也复制)。
registry 解析路径时优先用 `tools/bin/<tool>`
### 长期容器模式
daemon 启动时创建常驻容器 `ias-sandbox``tools/bin/` 只读挂载到 `/tools`
```
docker run -d --rm --name ias-sandbox \
--network none --cap-drop ALL --security-opt no-new-privileges \
--read-only --memory 512m --cpus 1 --pids-limit 128 \
--mount type=bind,source=tools/bin,target=/tools,readonly \
--mount type=tmpfs,target=/tmp \
ias-stdio-runner sleep infinity
```
工具调用通过 `docker exec` 在容器内执行(省去每次 `docker run` 的创建开销):
```
docker exec -i -e KEY=VAL ias-sandbox timeout -k 5 <secs> /tools/<name> --command <cmd>
```
- 容器已运行则复用,daemon 重启不重建
- `timeout -k 5` 超时 SIGTERM、5s 后 SIGKILL,避免孤儿进程
- env 按 spec 白名单 `-e` 注入(容器启动时仅有 PATH
- 新工具由 tool_manager 编译后复制到 `tools/bin/`,bind mount 自动可见,无需重启容器
### 两类基础镜像
| 镜像 | 用途 | 网络 | 文件系统 |
|------|------|------|----------|
| `ias-stdio-runner` | tool_manager 产出的不受信任工具(长期容器) | none | 只读根 + tmpfs /tmp |
| `ias-tool-builder` | tool_manager 编译 Rust 工具(一次性容器) | none | 可写 /workbind mount |
构建:`docker/build-images.sh`
### 容器安全基线
所有容器统一:
- `--rm` 一次性,用完销毁
- `--cap-drop ALL` 丢弃所有 Linux capability
- `--security-opt no-new-privileges` 禁止提权
- `--memory` / `--cpus` / `--pids-limit` 资源限制
- 只挂载工具二进制(readonly+ tmpfs /tmp
- 不挂载项目根目录、不挂载 docker.sock
- env 按白名单注入(`env_clear` + spec.env
### 沙箱配置(spec
```yaml
sandbox:
profile: stdio # local | stdio | builder
network: none # none | default
memory: 256m
cpus: "0.5"
pids_limit: 64
read_only: true
```
默认 `profile: local`(不启用容器)。开发者预写工具保持默认。
**tool_manager 产出的工具自动写入 `sandbox: profile: stdio`**(见 tool_manager `spec_yaml`),
不受信任代码默认进容器。
### 启用沙箱
```bash
export IAS_DOCKER_SANDBOX=1 # 工具执行走长期容器 docker exec
export IAS_DOCKER_BUILDER=1 # tool_manager 编译走 builder 容器
```
未设置时退回本地执行(仍保留 env_clear + 白名单注入)。
daemon 启动时自动 ensure `ias-sandbox` 容器,失败则回退本地执行。
### builder 容器特殊处理
`--cap-drop ALL` 丢掉 `CAP_DAC_OVERRIDE`root 无法写 bind mount 的非 root 目录。
解法:builder 容器以宿主 `uid:gid` 运行(`--user`),自然能写自己的目录。
镜像内 `CARGO_HOME=/tmp/cargo` 预缓存 serde_json 并 `chmod 777`,任意 uid 可用。
编译加 `--offline`,确保 `--network none` 下不尝试联网。
## tool_manager 加固
| 措施 | 实现 |
|------|------|
| 禁止自改 | `RESERVED_NAMES` 包含 `tool_manager`,禁止 create/update/build |
| 强制 high risk | `spec_yaml` 忽略调用方 risk_level,一律写 `high` |
| 审计日志 | 每次调用追加 `tools/.tool_audit.log`ts/uid/action/tool/result |
| 构建超时 | `timeout 180s` 包裹 cargo build |
| 失败清理 | 构建失败删除残留旧二进制,避免 reload 后调用旧版本"假成功" |
| 离线编译 | `IAS_DOCKER_BUILDER=1` 时在 builder 容器内 `--offline` 编译 |
## 结构化结果
`executor` 返回 `ToolRunOutcome` 枚举替代裸字符串:
```rust
enum ToolRunOutcome {
Ok(String),
Err { kind: ToolErrorKind, message: String },
}
```
`daemon` 据此判断是否 reloadtool_manager 成功后重载注册表),
替代原 `is_tool_error()` 的中文前缀字符串匹配。
## 审批并发修复
`ApprovalManager::clean_expired()` 两阶段加锁:
收集过期 `code_hash` → 更新 DB → 重新加锁按 `code_hash` 定位删除。
不依赖快照 index,避免并发回复导致 index 错位误删未过期审批。
## 验证
| 场景 | 结果 |
|------|------|
| tool_manager update 自身 | ❌ 拒绝(保留名) |
| create_tool risk_level=low | spec 强制 high |
| tool_manager 产出工具 spec | 自动含 `sandbox: profile: stdio` + 产物复制到 `tools/bin/` |
| datetime 本地执行 | +08:00(宿主时区,开发者工具默认不套容器) |
| tool_manager 产出工具 docker exec | 容器内 hostname=容器ID,证明确实走长期容器 exec |
| 容器常驻 | docker ps 可见 ias-sandboxexec 完无孤儿进程 |
| builder 容器编译 serde_json | ✅ 离线成功 |
| cargo test | 25 passed |
## 实施状态
| # | 步骤 | 状态 |
|---|------|------|
| 1 | 统一工具级审批 | ✅ |
| 2 | 禁止 tool_manager 自改 | ✅ |
| 3 | 强制生成工具 high risk | ✅ |
| 4 | env_clear 环境白名单 | ✅ |
| 5 | 修 clean_expired 并发 | ✅ |
| 6 | 三类 Docker 镜像 + sandbox spec | ✅ |
| 7 | stdio 工具迁移到一次性容器 | ✅ |
| 8 | network 工具迁移 | ➖ 不需要:工具走宿主代理联网,子进程纯计算,套容器无隔离收益 |
| 9 | tool_manager build 迁移到 builder 容器 | ✅ |
| 10 | 结构化结果替换 is_tool_error | ✅ |
@@ -0,0 +1,4 @@
-- 为"按用户查询最近的会话边界(clear/timeout 摘要)"添加索引。
-- list_recent_chat_records_for_context 会据此过滤掉已过期会话的消息。
CREATE INDEX IF NOT EXISTS idx_session_summaries_user_created
ON session_summaries (user_id, created_at DESC);
+77 -3
View File
@@ -260,13 +260,48 @@ pub async fn trigger_overflow_summary(
/// 而是只保存到数据库,供 `read_summaries` 工具查询。 /// 而是只保存到数据库,供 `read_summaries` 工具查询。
/// ///
/// 生成摘要后会清空所有消息,checkpoint 重置为 0。 /// 生成摘要后会清空所有消息,checkpoint 重置为 0。
/// 该摘要的 `created_at` 同时作为"会话边界"——之后加载历史时
/// 不会包含此时间点之前的消息(见 `db::latest_session_boundary`)。
pub async fn trigger_idle_summary( pub async fn trigger_idle_summary(
session: &Arc<Mutex<ChatSession>>, session: &Arc<Mutex<ChatSession>>,
summarizer: Option<&Summarizer>, summarizer: Option<&Summarizer>,
) { ) {
let _ = archive_session(session, summarizer, super::types::SummaryReason::Timeout, "timeout")
.await;
}
/// ## 触发清空上下文摘要(`/clear` 指令)
///
/// 用户发送 `/clear` 时调用。把当前会话消息交给 LLM 生成摘要并存档,
/// 然后清空消息、重置 checkpoint,使会话"过期"。
///
/// 与空闲超时摘要一样:
/// * 摘要**不会**注入到下一次对话,只存档供 `read_summaries` 查询
/// * 摘要的 `created_at` 作为会话边界,之后加载历史不再包含此前的消息
///
/// 返回生成的摘要文本;若会话本就没有消息,返回 `None`。
pub async fn trigger_clear_summary(
session: &Arc<Mutex<ChatSession>>,
summarizer: Option<&Summarizer>,
) -> Option<String> {
archive_session(session, summarizer, super::types::SummaryReason::Clear, "clear").await
}
/// 归档当前会话:生成摘要 → 存库 → 清空消息 → 重置 checkpoint。
///
/// `trigger_idle_summary` 与 `trigger_clear_summary` 共用此逻辑,
/// 仅 `reason` 不同(Timeout / Clear),二者都代表"会话过期"。
///
/// 返回生成的摘要文本;若会话没有消息可归档,返回 `None`。
async fn archive_session(
session: &Arc<Mutex<ChatSession>>,
summarizer: Option<&Summarizer>,
reason: super::types::SummaryReason,
reason_str: &str,
) -> Option<String> {
let mut s = session.lock().await; let mut s = session.lock().await;
if s.messages.is_empty() { if s.messages.is_empty() {
return; return None;
} }
let summary_text = generate_summary(&s.messages, summarizer).await; let summary_text = generate_summary(&s.messages, summarizer).await;
@@ -275,12 +310,12 @@ pub async fn trigger_idle_summary(
s.summaries.push(super::types::SummaryEntry { s.summaries.push(super::types::SummaryEntry {
checkpoint: cp, checkpoint: cp,
text: summary_text.clone(), text: summary_text.clone(),
reason: super::types::SummaryReason::Timeout, reason,
created_at: chrono::Utc::now(), created_at: chrono::Utc::now(),
}); });
let msg_count = cp as i32; let msg_count = cp as i32;
s.save_summary_to_db(&summary_text, "timeout", msg_count) s.save_summary_to_db(&summary_text, reason_str, msg_count)
.await; .await;
s.checkpoint = s.messages.len(); s.checkpoint = s.messages.len();
@@ -296,6 +331,8 @@ pub async fn trigger_idle_summary(
} }
} }
} }
Some(summary_text)
} }
/// 生成摘要:优先使用 LLM,不可用时回退到简单截断 /// 生成摘要:优先使用 LLM,不可用时回退到简单截断
@@ -431,4 +468,41 @@ mod tests {
assert!(has_anchor); assert!(has_anchor);
assert!(has_preface); assert!(has_preface);
} }
#[tokio::test]
async fn trigger_clear_summary_archives_and_clears() {
use crate::context::types::SummaryReason;
let mut session = ChatSession::new();
session.add_user("今天天气怎么样".to_string());
session.add_assistant("北京今天晴,25度".to_string());
session.add_user("明天呢".to_string());
session.add_assistant("明天多云转小雨".to_string());
let arc = Arc::new(Mutex::new(session));
let summarizer: Summarizer = Arc::new(|_prompt: String| {
Box::pin(async { Ok("摘要:讨论了北京近两天天气".to_string()) })
});
let summary = trigger_clear_summary(&arc, Some(&summarizer)).await;
assert!(summary.is_some());
assert_eq!(summary.unwrap(), "摘要:讨论了北京近两天天气");
let s = arc.lock().await;
assert!(s.messages.is_empty(), "清空后消息应被清空");
assert_eq!(s.checkpoint, 0);
assert!(s
.summaries
.iter()
.any(|x| x.reason == SummaryReason::Clear));
}
#[tokio::test]
async fn trigger_clear_summary_empty_session_returns_none() {
let arc = Arc::new(Mutex::new(ChatSession::new()));
let summarizer: Summarizer =
Arc::new(|_p: String| Box::pin(async { Ok("x".to_string()) }));
let summary = trigger_clear_summary(&arc, Some(&summarizer)).await;
assert!(summary.is_none());
}
} }
+6 -1
View File
@@ -14,6 +14,7 @@
//! 2. AI 回复 → `add_assistant()` → 存入 `messages` //! 2. AI 回复 → `add_assistant()` → 存入 `messages`
//! 3. Token 超 budget → `trigger_overflow_summary()` → 压缩 checkpoint 前的消息 //! 3. Token 超 budget → `trigger_overflow_summary()` → 压缩 checkpoint 前的消息
//! 4. 12h 空闲 → `trigger_idle_summary()` → 生成摘要,清空消息 //! 4. 12h 空闲 → `trigger_idle_summary()` → 生成摘要,清空消息
//! 5. `/clear` 指令 → `trigger_clear_summary()` → 生成摘要归档,会话过期
use crate::llm::types::Message; use crate::llm::types::Message;
use chrono::{DateTime, Utc}; use chrono::{DateTime, Utc};
@@ -30,15 +31,18 @@ pub struct HistoryEntry {
/// ## 摘要原因 /// ## 摘要原因
/// ///
/// 区分两种触发摘要的场景: /// 区分触发摘要的场景:
/// * `Overflow` — 上下文 token 总数超过预算(token budget),触发压缩 /// * `Overflow` — 上下文 token 总数超过预算(token budget),触发压缩
/// * `Timeout` — 距上条用户消息超过 12 小时空闲,触发会话总结 /// * `Timeout` — 距上条用户消息超过 12 小时空闲,触发会话总结
/// * `Clear` — 用户发送 `/clear` 主动清空上下文,触发会话过期归档
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum SummaryReason { pub enum SummaryReason {
/// 上下文 token 超 budget 触发 /// 上下文 token 超 budget 触发
Overflow, Overflow,
/// 距上条用户消息超过 12 小时触发 /// 距上条用户消息超过 12 小时触发
Timeout, Timeout,
/// 用户发送 `/clear` 主动清空上下文触发
Clear,
} }
/// ## 一条摘要记录 /// ## 一条摘要记录
@@ -77,6 +81,7 @@ pub struct SummaryEntry {
/// 2. AI 回复 → `add_assistant()` → 存入 `messages` /// 2. AI 回复 → `add_assistant()` → 存入 `messages`
/// 3. Token 超 budget → `trigger_overflow_summary()` → 压缩 checkpoint 前的消息 /// 3. Token 超 budget → `trigger_overflow_summary()` → 压缩 checkpoint 前的消息
/// 4. 12h 空闲 → `trigger_idle_summary()` → 生成摘要,清空消息 /// 4. 12h 空闲 → `trigger_idle_summary()` → 生成摘要,清空消息
/// 5. `/clear` 指令 → `trigger_clear_summary()` → 生成摘要归档,会话过期
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub struct ChatSession { pub struct ChatSession {
pub user_id: String, pub user_id: String,
+204 -38
View File
@@ -39,6 +39,7 @@
//! 采用统一的单进程 tokio task 架构。旧版 daemon+worker 分离模式(UDS IPC)已在 v2 中移除。 //! 采用统一的单进程 tokio task 架构。旧版 daemon+worker 分离模式(UDS IPC)已在 v2 中移除。
use crate::channel::ChannelId; use crate::channel::ChannelId;
use crate::context::builder;
use crate::context::HistoryEntry; use crate::context::HistoryEntry;
use crate::context::MemoryStore; use crate::context::MemoryStore;
use crate::db::Database; use crate::db::Database;
@@ -49,6 +50,7 @@ use crate::llm::{
use crate::queue::{ConsumerChannels, EnqueueHandle, MessageKind, PipelineMessage, QueueRunner}; use crate::queue::{ConsumerChannels, EnqueueHandle, MessageKind, PipelineMessage, QueueRunner};
use crate::tools::approval::ApprovalManager; use crate::tools::approval::ApprovalManager;
use crate::tools::executor::ToolRunOutcome;
use crate::tools::registry::RegistryManager; use crate::tools::registry::RegistryManager;
use crate::wechat::client::WeChatClient; use crate::wechat::client::WeChatClient;
use std::collections::HashMap; use std::collections::HashMap;
@@ -58,6 +60,10 @@ use tokio::sync::Mutex;
use tracing::{error, info, warn}; use tracing::{error, info, warn};
use uuid::Uuid; use uuid::Uuid;
/// LLM Consumer 的会话缓存:`correlation_id → (Conversation, 创建时间, user_id)`。
/// 额外记录 `user_id` 是为了在 `/clear` 时能按用户丢弃尚未完成的在途会话。
type SessionMap = HashMap<Uuid, (Conversation, Instant, String)>;
/// ## Daemon 上下文 —— 消费者之间共享的全部状态 /// ## Daemon 上下文 —— 消费者之间共享的全部状态
/// ///
/// 这个结构体被所有三个消费者(LLM / Tool / Send)通过 `Arc` 共享访问。 /// 这个结构体被所有三个消费者(LLM / Tool / Send)通过 `Arc` 共享访问。
@@ -111,6 +117,15 @@ pub async fn run(database: &Arc<Database>, memory_store: &Arc<MemoryStore>) {
database.pool().clone(), database.pool().clone(),
)))); ))));
let registry = Arc::new(RegistryManager::load("tools")); let registry = Arc::new(RegistryManager::load("tools"));
// 沙箱容器(如果启用):daemon 启动时创建长期容器,失败则设置降级标志
// executor 会据此回退本地执行,避免不受信任工具持续 SpawnFailed
if crate::tools::sandbox::enabled() {
if let Err(e) = crate::tools::sandbox::ensure_container("tools").await {
tracing::warn!(target: "ias::sandbox", "沙箱容器启动失败: {e}");
}
}
let system_prompt = std::env::var("WEIXIN_LLM_SYSTEM_PROMPT").unwrap_or_else(|_| String::new()); let system_prompt = std::env::var("WEIXIN_LLM_SYSTEM_PROMPT").unwrap_or_else(|_| String::new());
let model = std::env::var("DEEPSEEK_MODEL").unwrap_or_else(|_| "deepseek-v4-flash".to_string()); let model = std::env::var("DEEPSEEK_MODEL").unwrap_or_else(|_| "deepseek-v4-flash".to_string());
@@ -184,7 +199,7 @@ pub async fn run(database: &Arc<Database>, memory_store: &Arc<MemoryStore>) {
sessions_for_clean sessions_for_clean
.lock() .lock()
.await .await
.retain(|_, (conv, created)| { .retain(|_, (conv, created, _)| {
let ttl_secs = if conv.has_pending_async() { 1800 } else { 300 }; let ttl_secs = if conv.has_pending_async() { 1800 } else { 300 };
now.duration_since(*created).as_secs() < ttl_secs now.duration_since(*created).as_secs() < ttl_secs
}); });
@@ -260,9 +275,9 @@ fn spawn_consumers(
enqueue: EnqueueHandle, enqueue: EnqueueHandle,
) -> ( ) -> (
Vec<tokio::task::JoinHandle<()>>, Vec<tokio::task::JoinHandle<()>>,
Arc<Mutex<HashMap<Uuid, (Conversation, Instant)>>>, Arc<Mutex<SessionMap>>,
) { ) {
let sessions: Arc<Mutex<HashMap<Uuid, (Conversation, Instant)>>> = let sessions: Arc<Mutex<SessionMap>> =
Arc::new(Mutex::new(HashMap::new())); Arc::new(Mutex::new(HashMap::new()));
let mut handles = Vec::new(); let mut handles = Vec::new();
@@ -302,7 +317,7 @@ async fn llm_consumer_loop(
rx: &mut tokio::sync::mpsc::Receiver<PipelineMessage>, rx: &mut tokio::sync::mpsc::Receiver<PipelineMessage>,
ctx: &DaemonCtx, ctx: &DaemonCtx,
enqueue: &EnqueueHandle, enqueue: &EnqueueHandle,
sessions: &Arc<Mutex<HashMap<Uuid, (Conversation, Instant)>>>, sessions: &Arc<Mutex<SessionMap>>,
) { ) {
while let Some(msg) = rx.recv().await { while let Some(msg) = rx.recv().await {
let user_id = msg.channel.user_id.clone(); let user_id = msg.channel.user_id.clone();
@@ -314,6 +329,19 @@ async fn llm_consumer_loop(
message_id: _, message_id: _,
context_token, context_token,
} => { } => {
// /clear 指令:截断上下文 + 生成摘要 + 按会话过期处理
if is_clear_command(&text) {
handle_clear_command(
ctx,
enqueue,
sessions,
&user_id,
correlation_id,
context_token,
)
.await;
continue;
}
// 加载数据 // 加载数据
ctx.memory_store.load_if_needed(&user_id).await; ctx.memory_store.load_if_needed(&user_id).await;
let history = load_history(&ctx.db, &user_id).await; let history = load_history(&ctx.db, &user_id).await;
@@ -418,7 +446,7 @@ async fn llm_consumer_loop(
sessions sessions
.lock() .lock()
.await .await
.insert(correlation_id, (conv, Instant::now())); .insert(correlation_id, (conv, Instant::now(), user_id.clone()));
// 执行 LLM 对话 // 执行 LLM 对话
run_llm_round( run_llm_round(
@@ -442,7 +470,7 @@ async fn llm_consumer_loop(
// 从缓存恢复会话,注入工具结果并恢复 LLM 循环 // 从缓存恢复会话,注入工具结果并恢复 LLM 循环
let should_resume = { let should_resume = {
let mut map = sessions.lock().await; let mut map = sessions.lock().await;
if let Some((conv, last_access)) = map.get_mut(&correlation_id) { if let Some((conv, last_access, _)) = map.get_mut(&correlation_id) {
conv.session().lock().await.add_tool_result( conv.session().lock().await.add_tool_result(
&tool_call_id, &tool_call_id,
&tool_name, &tool_name,
@@ -492,13 +520,13 @@ async fn run_llm_round(
correlation_id: Uuid, correlation_id: Uuid,
context_token: Option<String>, context_token: Option<String>,
text: &str, text: &str,
sessions: &Arc<Mutex<HashMap<Uuid, (Conversation, Instant)>>>, sessions: &Arc<Mutex<SessionMap>>,
enqueue: &EnqueueHandle, enqueue: &EnqueueHandle,
) { ) {
let result = { let result = {
let map = sessions.lock().await; let map = sessions.lock().await;
let conv = match map.get(&correlation_id) { let conv = match map.get(&correlation_id) {
Some((c, _)) => c, Some((c, _, _)) => c,
None => { None => {
error!(target: "ias::err","[llm-session] conv不在缓存"); error!(target: "ias::err","[llm-session] conv不在缓存");
return; return;
@@ -535,13 +563,13 @@ async fn resume_llm_round(
user_id: &str, user_id: &str,
correlation_id: Uuid, correlation_id: Uuid,
context_token: Option<String>, context_token: Option<String>,
sessions: &Arc<Mutex<HashMap<Uuid, (Conversation, Instant)>>>, sessions: &Arc<Mutex<SessionMap>>,
enqueue: &EnqueueHandle, enqueue: &EnqueueHandle,
) { ) {
let result = { let result = {
let map = sessions.lock().await; let map = sessions.lock().await;
let conv = match map.get(&correlation_id) { let conv = match map.get(&correlation_id) {
Some((c, _)) => c, Some((c, _, _)) => c,
None => { None => {
error!(target: "ias::err","[llm-resume] conv不在缓存"); error!(target: "ias::err","[llm-resume] conv不在缓存");
return; return;
@@ -579,7 +607,7 @@ async fn handle_chat_result(
correlation_id: Uuid, correlation_id: Uuid,
context_token: Option<String>, context_token: Option<String>,
result: ChatResult, result: ChatResult,
sessions: &Arc<Mutex<HashMap<Uuid, (Conversation, Instant)>>>, sessions: &Arc<Mutex<SessionMap>>,
enqueue: &EnqueueHandle, enqueue: &EnqueueHandle,
) { ) {
if result.has_pending_async { if result.has_pending_async {
@@ -607,6 +635,111 @@ async fn handle_chat_result(
sessions.lock().await.remove(&correlation_id); sessions.lock().await.remove(&correlation_id);
} }
// ─── /clear 指令处理 ───
/// 判断文本是否为清空上下文指令 `/clear`(大小写不敏感,允许前后空白)。
fn is_clear_command(text: &str) -> bool {
text.trim().eq_ignore_ascii_case("/clear")
}
/// ## 处理 `/clear` 指令 —— 截断上下文 + 生成摘要 + 会话过期
///
/// 用户发送 `/clear` 时:
/// 1. 加载当前历史,用 LLM 生成会话摘要并归档(reason = clear
/// 2. 摘要的 `created_at` 成为会话边界,之后加载历史不再包含此前的消息
/// 3. 丢弃该用户尚未完成的在途会话(按会话过期处理)
/// 4. 回复用户清空确认(附带摘要预览)
async fn handle_clear_command(
ctx: &DaemonCtx,
enqueue: &EnqueueHandle,
sessions: &Arc<Mutex<SessionMap>>,
user_id: &str,
correlation_id: Uuid,
context_token: Option<String>,
) {
info!(target: "ias::msg", "[clear] {} 请求清空上下文", user_id);
// 1. 生成摘要:用一次性 Conversation 提供 summarizer + session
// 过滤掉历史里的 /clear 指令本身,避免污染摘要
let history: Vec<HistoryEntry> = load_history(&ctx.db, user_id)
.await
.into_iter()
.filter(|h| !is_clear_command(&h.content))
.collect();
let summary = if history.is_empty() {
None
} else {
let cfg = ConversationConfig {
system_prompt: String::new(),
model: ctx.model.clone(),
..Default::default()
};
match Conversation::new(cfg) {
Ok(conv) => {
{
let s = conv.session();
let mut s = s.lock().await;
s.db_pool = Some(Arc::new(ctx.db.pool().clone()));
s.load_from_history(&history);
s.current_user_id = user_id.to_string();
}
let summarizer = conv.summarizer();
builder::trigger_clear_summary(&conv.session(), Some(&summarizer)).await
}
Err(e) => {
error!(target: "ias::err", "[clear] 初始化摘要器失败: {e}");
None
}
}
};
// 2. 丢弃该用户在途的会话(按会话过期处理)
sessions
.lock()
.await
.retain(|_, (_, _, uid)| uid.as_str() != user_id);
// 3. 回复确认(附带摘要预览)
let reply = build_clear_reply(&summary);
let msg = PipelineMessage::llm_reply_with_source(
ChannelId::wechat(user_id),
correlation_id,
&reply,
"clear",
context_token,
);
enqueue.enqueue(msg).await;
}
/// 根据 `trigger_clear_summary` 的返回值构建 `/clear` 回复文案。
///
/// * `None` — 无历史可清空
/// * `Some("")` — 已清空但 LLM 摘要为空
/// * `Some(text)` — 已清空,附带 300 字预览
fn build_clear_reply(summary: &Option<String>) -> String {
match summary {
None => "🧹 上下文已是空的,无需清空。".to_string(),
Some(text) if text.is_empty() => "🧹 已清空上下文,本次对话已归档。".to_string(),
Some(text) => {
let chars: Vec<char> = text.chars().collect();
let (preview, suffix) = if chars.len() > 300 {
(
chars[..300].iter().collect::<String>(),
"\n…(摘要已截断预览,完整内容已归档)",
)
} else {
(text.clone(), "")
};
format!(
"🧹 已清空上下文,本次对话已归档为摘要。\n\n📝 摘要预览:\n{}{}",
preview, suffix
)
}
}
}
// ─── Tool Consumer ─── // ─── Tool Consumer ───
async fn tool_consumer_loop( async fn tool_consumer_loop(
@@ -692,7 +825,7 @@ async fn handle_tool_call(
eq.enqueue(msg).await; eq.enqueue(msg).await;
}) })
}; };
let output = ctx let outcome = ctx
.registry .registry
.dispatch( .dispatch(
target, target,
@@ -704,11 +837,17 @@ async fn handle_tool_call(
&wechat_cb, &wechat_cb,
) )
.await; .await;
if target == "tool_manager" && !is_tool_error(&output) { if target == "tool_manager" && !outcome.is_err() {
let count = ctx.registry.reload().await; let count = ctx.registry.reload().await;
format!("{output}\n\n工具注册表已重载,当前可用工具数: {count}") format!(
"{}\n\n工具注册表已重载,当前可用工具数: {count}",
outcome.into_message()
)
} else { } else {
output if let ToolRunOutcome::Err { ref kind, .. } = outcome {
tracing::warn!(target:"ias::err","[tool] {target} 执行失败: {kind:?}");
}
outcome.into_message()
} }
} else { } else {
let params: serde_json::Value = serde_json::from_str(&arguments).unwrap_or_default(); let params: serde_json::Value = serde_json::from_str(&arguments).unwrap_or_default();
@@ -734,7 +873,7 @@ async fn handle_tool_call(
eq.enqueue(msg).await; eq.enqueue(msg).await;
}) })
}; };
let output = ctx let outcome = ctx
.registry .registry
.dispatch( .dispatch(
&tool_name, &tool_name,
@@ -746,11 +885,17 @@ async fn handle_tool_call(
&wechat_cb, &wechat_cb,
) )
.await; .await;
if tool_name == "tool_manager" && !is_tool_error(&output) { if tool_name == "tool_manager" && !outcome.is_err() {
let count = ctx.registry.reload().await; let count = ctx.registry.reload().await;
format!("{output}\n\n工具注册表已重载,当前可用工具数: {count}") format!(
"{}\n\n工具注册表已重载,当前可用工具数: {count}",
outcome.into_message()
)
} else { } else {
output if let ToolRunOutcome::Err { ref kind, .. } = outcome {
tracing::warn!(target:"ias::err","[tool] {tool_name} 执行失败: {kind:?}");
}
outcome.into_message()
} }
}; };
@@ -767,25 +912,6 @@ async fn handle_tool_call(
.await; .await;
} }
fn is_tool_error(result: &str) -> bool {
result.starts_with("异常退出")
|| result.starts_with("启动失败")
|| result.starts_with("工具输出非JSON")
|| result.starts_with("HTTP失败")
|| result.starts_with("读取stdout失败")
|| result.starts_with("stdout超时")
|| result.starts_with("进程超时")
|| result.starts_with("进程异常")
|| result.starts_with("未知消息类型")
|| result.starts_with("未知工具")
|| result.starts_with("不支持的方法")
|| result.starts_with("审批失败")
|| result.starts_with("审批超时")
|| result.starts_with("操作已被用户取消")
|| result.starts_with("缺少必填参数")
|| result.starts_with("工具 '")
}
fn unpack_call_capability_params(cp: &serde_json::Value) -> (&str, serde_json::Value) { fn unpack_call_capability_params(cp: &serde_json::Value) -> (&str, serde_json::Value) {
let target = cp["name"].as_str().unwrap_or(""); let target = cp["name"].as_str().unwrap_or("");
let mut params = serde_json::Map::new(); let mut params = serde_json::Map::new();
@@ -961,3 +1087,43 @@ async fn load_summaries(db: &Arc<Database>, user_id: &str) -> Vec<String> {
} }
} }
} }
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn clear_command_detection() {
assert!(is_clear_command("/clear"));
assert!(is_clear_command(" /clear "));
assert!(is_clear_command("/CLEAR"));
assert!(is_clear_command("/Clear"));
assert!(is_clear_command("\t/clear\n"));
// 非精确匹配不应触发
assert!(!is_clear_command("/clear all"));
assert!(!is_clear_command("请/clear"));
assert!(!is_clear_command("清空"));
assert!(!is_clear_command(""));
assert!(!is_clear_command("/clea"));
}
#[test]
fn clear_reply_branches() {
// 无历史可清空
assert_eq!(build_clear_reply(&None), "🧹 上下文已是空的,无需清空。");
// 已清空但摘要为空(不应误报“已是空的”)
assert_eq!(
build_clear_reply(&Some(String::new())),
"🧹 已清空上下文,本次对话已归档。"
);
// 正常摘要:附带预览
let r = build_clear_reply(&Some("讨论了天气".to_string()));
assert!(r.starts_with("🧹 已清空上下文,本次对话已归档为摘要。"));
assert!(r.contains("讨论了天气"));
// 超长摘要:截断预览
let long = "x".repeat(500);
let r = build_clear_reply(&Some(long));
assert!(r.contains("摘要已截断预览"));
assert!(!r.contains(&"x".repeat(500)));
}
}
+37 -2
View File
@@ -138,22 +138,31 @@ pub async fn insert_chat_record(
} }
/// 查询最近可进入 LLM 上下文的聊天记录(排除工具等待/审批等中间提示)。 /// 查询最近可进入 LLM 上下文的聊天记录(排除工具等待/审批等中间提示)。
///
/// 会自动应用"会话边界":最近一次 `/clear` 或 12h 空闲超时摘要之前的消息
/// 已归档为摘要,不再进入新会话的上下文(见 [`latest_session_boundary`])。
pub async fn list_recent_chat_records_for_context( pub async fn list_recent_chat_records_for_context(
pool: &PgPool, pool: &PgPool,
user_id: &str, user_id: &str,
limit: i64, limit: i64,
) -> Result<Vec<ChatRecord>, String> { ) -> Result<Vec<ChatRecord>, String> {
// 会话边界:最近一次"会话过期"摘要(/clear 或 12h 空闲超时)的时间点。
// 边界之前的消息已归档为摘要,不再进入新会话的上下文。
let boundary = latest_session_boundary(pool, user_id).await;
let records = sqlx::query_as::<_, ChatRecord>( let records = sqlx::query_as::<_, ChatRecord>(
r#" r#"
SELECT id, created_at, direction, user_id, account_id, text, source, context_token, message_id SELECT id, created_at, direction, user_id, account_id, text, source, context_token, message_id
FROM chat_records FROM chat_records
WHERE user_id = $1 WHERE user_id = $1
AND source NOT IN ('llm_pending_tool', 'tool_approval') AND source NOT IN ('llm_pending_tool', 'tool_approval', 'clear')
AND created_at > COALESCE($2, '-infinity'::timestamptz)
ORDER BY created_at DESC ORDER BY created_at DESC
LIMIT $2 LIMIT $3
"#, "#,
) )
.bind(user_id) .bind(user_id)
.bind(boundary)
.bind(limit) .bind(limit)
.fetch_all(pool) .fetch_all(pool)
.await .await
@@ -162,6 +171,32 @@ pub async fn list_recent_chat_records_for_context(
Ok(records) Ok(records)
} }
/// ## 查询用户最近的"会话边界"时间点
///
/// 返回最近一次会话过期摘要(`/clear` 或 12h 空闲超时)的 `created_at`。
/// 该时间点之前的聊天记录已归档为摘要,不应再进入新会话的上下文。
///
/// `Overflow` 摘要不构成边界(它发生在会话进行中,仅压缩早期消息)。
pub async fn latest_session_boundary(
pool: &PgPool,
user_id: &str,
) -> Option<DateTime<Utc>> {
sqlx::query_as::<_, (DateTime<Utc>,)>(
r#"
SELECT created_at
FROM session_summaries
WHERE user_id = $1 AND reason IN ('clear', 'timeout')
ORDER BY created_at DESC
LIMIT 1
"#,
)
.bind(user_id)
.fetch_optional(pool)
.await
.ok()?
.map(|(t,)| t)
}
// ─── LLM 用量 ─── // ─── LLM 用量 ───
/// LLM 用量聚合统计 /// LLM 用量聚合统计
+4 -1
View File
@@ -478,7 +478,10 @@ pub const DEFAULT_SYSTEM_PROMPT: &str = "\
1. `query_capabilities` — 查询所有可用能力工具\n\ 1. `query_capabilities` — 查询所有可用能力工具\n\
2. `call_capability` — 调用能力工具,参数 {\"name\":\"工具名\",\"prompt\":{\"参数\":\"\"}}\n\ 2. `call_capability` — 调用能力工具,参数 {\"name\":\"工具名\",\"prompt\":{\"参数\":\"\"}}\n\
\n\ \n\
流程:收到消息 → query_capabilities → call_capability → 回复\n\ 流程:收到消息 → query_capabilities → call_capability → 回复
\
\
需要创建或修改本地工具时:先用 `tool_inspector`action=read_tool)查看现有实现,再用 `tool_manager` 修改。`tool_manager` 是高风险操作,执行前会请求用户确认。\n\
"; ";
#[cfg(test)] #[cfg(test)]
+7 -24
View File
@@ -529,7 +529,7 @@ async fn cmd_tool_dynamic(name: &str, args: &[String]) {
+ Send + Send
+ Sync + Sync
) = &|_: &str, _: &str| Box::pin(async {}); ) = &|_: &str, _: &str| Box::pin(async {});
let result = registry let outcome = registry
.dispatch( .dispatch(
name, name,
&serde_json::Value::Object(params), &serde_json::Value::Object(params),
@@ -541,28 +541,11 @@ async fn cmd_tool_dynamic(name: &str, args: &[String]) {
) )
.await; .await;
// 检测是否执行出错(匹配 execute_loop 返回的所有错误前缀) match outcome {
let is_error = result.starts_with("异常退出") crate::tools::executor::ToolRunOutcome::Ok(s) => println!("{s}"),
|| result.starts_with("启动失败") crate::tools::executor::ToolRunOutcome::Err { message, .. } => {
|| result.starts_with("工具输出非JSON") eprintln!("{message}");
|| result.starts_with("HTTP失败") std::process::exit(1);
|| result.starts_with("读取stdout失败") }
|| result.starts_with("stdout超时")
|| result.starts_with("进程超时")
|| result.starts_with("进程异常")
|| result.starts_with("未知消息类型")
|| result.starts_with("未知工具")
|| result.starts_with("不支持的方法")
|| result.starts_with("审批失败")
|| result.starts_with("审批超时")
|| result.starts_with("操作已被用户取消")
|| result.starts_with("缺少必填参数")
|| result.starts_with("工具 '");
if is_error {
eprintln!("{result}");
std::process::exit(1);
} else {
println!("{result}");
} }
} }
+113 -11
View File
@@ -129,6 +129,33 @@ impl ApprovalManager {
Ok((code_str, rx)) Ok((code_str, rx))
} }
/// 测试专用:创建指定 TTL(秒)的审批请求
#[cfg(test)]
async fn create_with_ttl(
&self,
user_id: &str,
skill_name: &str,
ttl_secs: u64,
) -> Result<(String, oneshot::Receiver<ApprovalDecision>), String> {
let code: u32 = rand::rng().random_range(100000..999999);
let code_str = code.to_string();
let code_hash = format!("{:x}", Sha256::digest(code_str.as_bytes()));
let (tx, rx) = oneshot::channel();
let mut map = self.pending.lock().await;
map.entry(user_id.to_string())
.or_default()
.push(PendingEntry {
code_hash,
skill_name: skill_name.to_string(),
attempts_left: 3,
expires_at: Instant::now() + std::time::Duration::from_secs(ttl_secs),
sender: tx,
});
Ok((code_str, rx))
}
/// 处理用户回复(由消息循环调用) /// 处理用户回复(由消息循环调用)
/// 返回 (技能名, 决策);匹配失败返回 None /// 返回 (技能名, 决策);匹配失败返回 None
pub async fn handle_reply( pub async fn handle_reply(
@@ -192,22 +219,32 @@ impl ApprovalManager {
} }
/// 清理过期审批(按时间判定,通知等待方 + 更新 DB) /// 清理过期审批(按时间判定,通知等待方 + 更新 DB)
///
/// 两阶段加锁:收集过期 code_hash → 更新 DB → 重新加锁按 code_hash 定位删除。
/// 不依赖快照 index,避免并发回复导致 index 错位误删。
pub async fn clean_expired(&self) { pub async fn clean_expired(&self) {
let now = Instant::now(); let now = Instant::now();
let mut expired = Vec::new();
// 第一阶段:持锁收集过期条目的 code_hash
let mut expired_hashes: Vec<String> = Vec::new();
{ {
let map = self.pending.lock().await; let map = self.pending.lock().await;
for (uid, entries) in map.iter() { for entries in map.values() {
for (i, e) in entries.iter().enumerate() { for e in entries {
if now > e.expires_at { if now > e.expires_at {
expired.push((uid.clone(), i, e.code_hash.clone())); expired_hashes.push(e.code_hash.clone());
} }
} }
} }
} }
if expired_hashes.is_empty() {
return;
}
// DB 更新(按 code_hash 精确匹配,无错位问题)
if let Some(ref pool) = self.pool { if let Some(ref pool) = self.pool {
for (_uid, _i, code_hash) in &expired { for code_hash in &expired_hashes {
let _ = sqlx::query( let _ = sqlx::query(
"UPDATE pending_approvals SET status = 'expired' WHERE code_hash = $1 AND status = 'pending'", "UPDATE pending_approvals SET status = 'expired' WHERE code_hash = $1 AND status = 'pending'",
) )
@@ -217,15 +254,80 @@ impl ApprovalManager {
} }
} }
// 第二阶段:重新持锁,按 code_hash 定位并移除
// (即使两阶段间 handle_reply 增删了条目,按 hash 查找也能安全跳过)
let mut map = self.pending.lock().await; let mut map = self.pending.lock().await;
for (uid, i, _) in expired.iter().rev() { for code_hash in &expired_hashes {
if let Some(entries) = map.get_mut(uid) for entries in map.values_mut() {
&& *i < entries.len() if let Some(idx) = entries.iter().position(|e| &e.code_hash == code_hash) {
{ let entry = entries.remove(idx);
let entry = entries.remove(*i); let _ = entry.sender.send(ApprovalDecision::Expired);
let _ = entry.sender.send(ApprovalDecision::Expired); }
} }
} }
map.retain(|_, v| !v.is_empty()); map.retain(|_, v| !v.is_empty());
} }
} }
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn clean_expired_only_removes_expired_entries() {
// 无 DB,纯内存
let mgr = ApprovalManager::new(None);
// 两条审批:一条立即过期(ttl=0),一条 300s 不过期
let (_expired_code, expired_rx) = mgr.create_with_ttl("u1", "tool_a", 0).await.unwrap();
let (live_code, live_rx) = mgr.create_with_ttl("u1", "tool_b", 300).await.unwrap();
// 立即过期的条目 expires_at 已成过去,sleep 确保时间推进
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
mgr.clean_expired().await;
// 过期条目应收到 Expired
assert_eq!(expired_rx.await.unwrap(), ApprovalDecision::Expired);
// 未过期条目不应被误删:handle_reply 用正确确认码应能匹配
let (name, decision) = mgr.handle_reply("u1", &live_code).await.unwrap();
assert_eq!(name, "tool_b");
assert_eq!(decision, ApprovalDecision::Approved);
// live_rx 应收到 Approved(未被 clean_expired 误标记)
assert_eq!(live_rx.await.unwrap(), ApprovalDecision::Approved);
}
#[tokio::test]
async fn clean_expired_safe_when_no_pending() {
// 空 pending 不应 panic
let mgr = ApprovalManager::new(None);
mgr.clean_expired().await;
// 重复调用也安全
mgr.clean_expired().await;
}
#[tokio::test]
async fn clean_expired_skips_entries_already_removed_by_reply() {
// 回归保护:clean_expired 第一阶段收集过期 hash 后、第二阶段删除前,
// 若该条目已被 handle_reply 提前处理,clean_expired 应安全跳过,不误删其他条目。
// (原 bug:按快照 index 删除会导致误删;现按 code_hash 删除应安全跳过)
let mgr = ApprovalManager::new(None);
// 两条同用户审批,都设为立即过期
let (_code_a, rx_a) = mgr.create_with_ttl("u1", "tool_a", 0).await.unwrap();
let (_code_b, rx_b) = mgr.create_with_ttl("u1", "tool_b", 0).await.unwrap();
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
// 先用 handle_reply “0” 取消最早的条目(tool_a),模拟并发回复
let (name, decision) = mgr.handle_reply("u1", "0").await.unwrap();
assert_eq!(name, "tool_a");
assert_eq!(decision, ApprovalDecision::Rejected);
assert_eq!(rx_a.await.unwrap(), ApprovalDecision::Rejected);
// 再 clean_expiredtool_a 已被 remove,按 hash 查找应跳过;
// tool_b 仍在且过期,应被正确标记 Expired(不会被误当成 tool_a 跳过)
mgr.clean_expired().await;
assert_eq!(rx_b.await.unwrap(), ApprovalDecision::Expired);
}
}
+193 -78
View File
@@ -5,7 +5,11 @@
//! - `{"type":"db",...}` — 请求数据库操作 //! - `{"type":"db",...}` — 请求数据库操作
//! - `{"type":"result","content":"..."}` — 最终结果 //! - `{"type":"result","content":"..."}` — 最终结果
//! //!
//! 所有消息必须是合法 JSON。审批在 http/db 操作前处理。 //! 所有消息必须是合法 JSON。
//!
//! ## 安全边界(授权层)
//! **工具级审批在入口统一执行**:高风险工具在 spawn 之前必须获得用户确认,
//! 与工具后续输出的消息类型(result/http/db)无关 —— 杜绝 stdio→result 绕过。
use std::sync::Arc; use std::sync::Arc;
use std::time::Duration; use std::time::Duration;
@@ -24,6 +28,59 @@ struct HttpResponse {
body: Value, body: Value,
} }
// ─── 结构化结果 ───
/// 工具执行错误类别(供调用方区分,替代字符串前缀匹配)
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ToolErrorKind {
SpawnFailed,
StdoutReadFailed,
StdoutTimeout,
ProcessTimeout,
ProcessAbnormal,
NonJsonOutput,
UnknownMessageType,
UnsupportedMethod,
HttpFailed,
ApprovalFailed,
ApprovalTimeout,
ApprovalRejected,
MaxRoundsExceeded,
UnknownTool,
MissingParams,
}
/// 工具执行返回值:成功带内容,失败带类别与消息
#[derive(Debug, Clone)]
pub enum ToolRunOutcome {
Ok(String),
Err {
kind: ToolErrorKind,
message: String,
},
}
impl ToolRunOutcome {
pub fn err(kind: ToolErrorKind, message: impl Into<String>) -> Self {
ToolRunOutcome::Err {
kind,
message: message.into(),
}
}
pub fn is_err(&self) -> bool {
matches!(self, ToolRunOutcome::Err { .. })
}
/// 取出最终展示给 LLM 的消息文本
pub fn into_message(self) -> String {
match self {
ToolRunOutcome::Ok(s) => s,
ToolRunOutcome::Err { message, .. } => message,
}
}
}
pub async fn execute_loop( pub async fn execute_loop(
spec: &ToolSpec, spec: &ToolSpec,
params: &Value, params: &Value,
@@ -36,11 +93,32 @@ pub async fn execute_loop(
+ Send + Send
+ Sync + Sync
), ),
) -> String { ) -> ToolRunOutcome {
// ── 统一授权边界:高风险工具在 spawn 前必须通过用户审批 ──
// 不依赖后续消息类型,覆盖 stdio/http/db 所有路径
if spec.is_high_risk() && !approved {
let (code, rx) = match approval.create(user_id, &spec.name).await {
Ok((c, r)) => (c, r),
Err(e) => return ToolRunOutcome::err(ToolErrorKind::ApprovalFailed, format!("审批失败: {e}")),
};
wechat_send(
user_id,
&format!("⚠️ {}\n确认码: {}\n回复确认码继续,回复0取消", spec.desc, code),
)
.await;
match rx.await.unwrap_or(ApprovalDecision::Expired) {
ApprovalDecision::Approved => { /* 授权通过,继续执行 */ }
ApprovalDecision::Rejected => {
return ToolRunOutcome::err(ToolErrorKind::ApprovalRejected, "操作已被用户取消")
}
ApprovalDecision::Expired => {
return ToolRunOutcome::err(ToolErrorKind::ApprovalTimeout, "审批超时,操作已取消")
}
}
}
let start = std::time::Instant::now(); let start = std::time::Instant::now();
let binary_path = spec.path.as_deref().unwrap_or(&spec.tool); let binary_path = spec.path.as_deref().unwrap_or(&spec.tool);
const MAX_ROUNDS: u32 = 20; const MAX_ROUNDS: u32 = 20;
let mut response: Option<HttpResponse> = None; let mut response: Option<HttpResponse> = None;
@@ -51,7 +129,10 @@ pub async fn execute_loop(
round += 1; round += 1;
if round > MAX_ROUNDS { if round > MAX_ROUNDS {
error!(target: "ias::err", "[loop] 工具 '{}' 超过最大轮次 {}", spec.name, MAX_ROUNDS); error!(target: "ias::err", "[loop] 工具 '{}' 超过最大轮次 {}", spec.name, MAX_ROUNDS);
return format!("工具 '{}' 执行轮次过多(>{})", spec.name, MAX_ROUNDS); return ToolRunOutcome::err(
ToolErrorKind::MaxRoundsExceeded,
format!("工具 '{}' 执行轮次过多(>{})", spec.name, MAX_ROUNDS),
);
} }
// ── 信封 ── // ── 信封 ──
@@ -64,32 +145,20 @@ pub async fn execute_loop(
} }
// ── spawn ── // ── spawn ──
let binary = std::path::Path::new(binary_path); // 执行隔离边界:启用沙箱且工具 profile 非 local 时,通过 docker exec
let mut cmd = Command::new(&binary); // 在长期容器内执行;否则本地 spawn(env_clear + 白名单注入)。
if let Some(c) = &spec.command { // 沙箱降级(ensure 失败或容器崩溃未恢复)时自动回退本地执行。
cmd.arg("--command").arg(c); let use_docker = if crate::tools::sandbox::should_use_docker(&spec.sandbox.profile) {
} // 执行前确保容器就绪:未运行则重建,重建失败则降级
// 按工具 spec 定义的 env 列表,从当前进程环境推送对应变量 crate::tools::sandbox::ensure_ready_or_degrade("tools").await
for key in &spec.env { } else {
match std::env::var(key) { false
Ok(v) => { };
cmd.env(key, &v); let mut child = match spawn_child(spec, binary_path, use_docker).await {
}
Err(_) => {
warn!(target: "ias::err", "[env] 工具 '{}' 需要的环境变量 '{}' 未设置", spec.name, key);
}
}
}
cmd.stdin(std::process::Stdio::piped())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.kill_on_drop(true);
let mut child = match cmd.spawn() {
Ok(c) => c, Ok(c) => c,
Err(e) => { Err(e) => {
error!(target: "ias::err","[spawn] {e}"); error!(target: "ias::err","[spawn] {e}");
return format!("启动失败: {e}"); return ToolRunOutcome::err(ToolErrorKind::SpawnFailed, format!("启动失败: {e}"));
} }
}; };
@@ -129,11 +198,14 @@ pub async fn execute_loop(
Ok(Ok(l)) => l, Ok(Ok(l)) => l,
Ok(Err(e)) => { Ok(Err(e)) => {
let _ = child.start_kill(); let _ = child.start_kill();
return format!("读取stdout失败: {e}"); return ToolRunOutcome::err(ToolErrorKind::StdoutReadFailed, format!("读取stdout失败: {e}"));
} }
Err(_) => { Err(_) => {
let _ = child.start_kill(); let _ = child.start_kill();
return format!("stdout超时({}s)", spec.timeout_secs); return ToolRunOutcome::err(
ToolErrorKind::StdoutTimeout,
format!("stdout超时({}s)", spec.timeout_secs),
);
} }
}; };
@@ -152,12 +224,15 @@ pub async fn execute_loop(
Ok(Ok(s)) => s, Ok(Ok(s)) => s,
Ok(Err(e)) => { Ok(Err(e)) => {
error!(target: "ias::err","[process] {e}"); error!(target: "ias::err","[process] {e}");
return format!("进程异常: {e}"); return ToolRunOutcome::err(ToolErrorKind::ProcessAbnormal, format!("进程异常: {e}"));
} }
Err(_) => { Err(_) => {
let _ = child.start_kill(); let _ = child.start_kill();
error!(target: "ias::err","[timeout] {}s", spec.timeout_secs); error!(target: "ias::err","[timeout] {}s", spec.timeout_secs);
return format!("进程超时({}s)", spec.timeout_secs); return ToolRunOutcome::err(
ToolErrorKind::ProcessTimeout,
format!("进程超时({}s)", spec.timeout_secs),
);
} }
}; };
@@ -168,7 +243,10 @@ pub async fn execute_loop(
} else { } else {
stderr_trimmed.to_string() stderr_trimmed.to_string()
}; };
return format!("异常退出({}): {}", status.code().unwrap_or(-1), reason); return ToolRunOutcome::err(
ToolErrorKind::ProcessAbnormal,
format!("异常退出({}): {}", status.code().unwrap_or(-1), reason),
);
} }
// ── 解析 ── // ── 解析 ──
@@ -177,7 +255,10 @@ pub async fn execute_loop(
Ok(v) => v, Ok(v) => v,
Err(_) => { Err(_) => {
warn!(target: "ias::err","[non-json] {:.100}", trimmed); warn!(target: "ias::err","[non-json] {:.100}", trimmed);
return format!("工具输出非JSON: {:.100}", trimmed); return ToolRunOutcome::err(
ToolErrorKind::NonJsonOutput,
format!("工具输出非JSON: {:.100}", trimmed),
);
} }
}; };
@@ -186,30 +267,11 @@ pub async fn execute_loop(
let content = msg["content"].as_str().unwrap_or(""); let content = msg["content"].as_str().unwrap_or("");
let elapsed = start.elapsed(); let elapsed = start.elapsed();
info!(target: "ias::tool", "[done] {} 完成, 耗时 {:?}, {} 轮", spec.name, elapsed, round); info!(target: "ias::tool", "[done] {} 完成, 耗时 {:?}, {} 轮", spec.name, elapsed, round);
return content.to_string(); return ToolRunOutcome::Ok(content.to_string());
} }
Some("http") => { Some("http") => {
let method = msg["method"].as_str().unwrap_or("GET"); let method = msg["method"].as_str().unwrap_or("GET");
let url = msg["url"].as_str().unwrap_or(""); let url = msg["url"].as_str().unwrap_or("");
let desc = msg["desc"].as_str().unwrap_or("");
if spec.is_high_risk() && !approved {
let (code, rx) = match approval.create(user_id, &spec.name).await {
Ok((c, r)) => (c, r),
Err(e) => return format!("审批失败: {e}"),
};
wechat_send(
user_id,
&format!("⚠️ {}\n确认码: {}\n回复确认码继续,回复0取消", desc, code),
)
.await;
let d = rx.await.unwrap_or(ApprovalDecision::Expired);
match d {
ApprovalDecision::Approved => { /* 继续执行 */ }
ApprovalDecision::Rejected => return "操作已被用户取消".to_string(),
ApprovalDecision::Expired => return "审批超时,操作已取消".to_string(),
}
}
let client = reqwest::Client::builder() let client = reqwest::Client::builder()
.timeout(Duration::from_secs(25)) .timeout(Duration::from_secs(25))
@@ -227,7 +289,9 @@ pub async fn execute_loop(
.await .await
} }
("POST", None) => client.post(url).send().await, ("POST", None) => client.post(url).send().await,
(o, _) => return format!("不支持的方法: {o}"), (o, _) => {
return ToolRunOutcome::err(ToolErrorKind::UnsupportedMethod, format!("不支持的方法: {o}"))
}
}; };
match resp { match resp {
Ok(r) => { Ok(r) => {
@@ -237,32 +301,13 @@ pub async fn execute_loop(
} }
Err(e) => { Err(e) => {
error!(target: "ias::err","[http] {e}"); error!(target: "ias::err","[http] {e}");
return format!("HTTP失败: {e}"); return ToolRunOutcome::err(ToolErrorKind::HttpFailed, format!("HTTP失败: {e}"));
} }
} }
} }
Some("db") => { Some("db") => {
let op = msg["operation"].as_str().unwrap_or(""); let op = msg["operation"].as_str().unwrap_or("");
let op_params = msg.get("params").cloned().unwrap_or(Value::Null); let op_params = msg.get("params").cloned().unwrap_or(Value::Null);
let desc = msg["desc"].as_str().unwrap_or(op);
if spec.is_high_risk() && !approved {
let (code, rx) = match approval.create(user_id, &spec.name).await {
Ok((c, r)) => (c, r),
Err(e) => return format!("审批失败: {e}"),
};
wechat_send(
user_id,
&format!("⚠️ {}\n确认码: {}\n回复确认码继续,回复0取消", desc, code),
)
.await;
let d = rx.await.unwrap_or(ApprovalDecision::Expired);
match d {
ApprovalDecision::Approved => { /* 继续执行 */ }
ApprovalDecision::Rejected => return "操作已被用户取消".to_string(),
ApprovalDecision::Expired => return "审批超时,操作已取消".to_string(),
}
}
let result = match op { let result = match op {
"read_memories" => memory.read_for(user_id, None).await, "read_memories" => memory.read_for(user_id, None).await,
@@ -271,8 +316,7 @@ pub async fn execute_loop(
if c.is_empty() { if c.is_empty() {
"缺少 content".into() "缺少 content".into()
} else { } else {
let write_result = memory.write_for(user_id, c).await; memory.write_for(user_id, c).await
write_result
} }
} }
"delete_memory" => { "delete_memory" => {
@@ -290,8 +334,79 @@ pub async fn execute_loop(
} }
_ => { _ => {
warn!(target: "ias::err","[unknown-type] {}", msg["type"]); warn!(target: "ias::err","[unknown-type] {}", msg["type"]);
return format!("未知消息类型: {}", msg["type"]); return ToolRunOutcome::err(
ToolErrorKind::UnknownMessageType,
format!("未知消息类型: {}", msg["type"]),
);
} }
} }
} }
} }
/// 本地执行:env_clear + PATH + spec.env 白名单注入
fn build_local_cmd(spec: &ToolSpec, binary_path: &str) -> Command {
let mut cmd = Command::new(std::path::Path::new(binary_path));
if let Some(c) = &spec.command {
cmd.arg("--command").arg(c);
}
cmd.env_clear();
if let Ok(path) = std::env::var("PATH") {
cmd.env("PATH", path);
}
for key in &spec.env {
match std::env::var(key) {
Ok(v) => {
cmd.env(key, &v);
}
Err(_) => {
warn!(target: "ias::err", "[env] 工具 '{}' 需要的环境变量 '{}' 未设置", spec.name, key);
}
}
}
cmd
}
/// spawn 工具进程。docker exec 模式下若首次 spawn 失败,尝试检测容器崩溃并重建后重试一次;
/// 重建失败则降级为本地执行(保留可用性,但仍受审批约束)。
async fn spawn_child(
spec: &ToolSpec,
binary_path: &str,
use_docker: bool,
) -> Result<tokio::process::Child, std::io::Error> {
let make_cmd = || {
let mut cmd = if use_docker {
crate::tools::sandbox::build_exec_cmd(spec)
} else {
build_local_cmd(spec, binary_path)
};
cmd.stdin(std::process::Stdio::piped())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.kill_on_drop(true);
cmd
};
if let Ok(child) = make_cmd().spawn() {
return Ok(child);
}
// 首次 spawn 失败。docker 模式下可能是容器崩溃,尝试重建后重试一次
if use_docker && crate::tools::sandbox::recover_if_crashed("tools").await {
tracing::info!(target: "ias::sandbox", "沙箱容器已重建,重试执行");
if let Ok(child) = make_cmd().spawn() {
return Ok(child);
}
// 重建后仍失败,降级本地执行
tracing::warn!(target: "ias::sandbox", "重建后 exec 仍失败,降级本地执行");
} else if use_docker {
tracing::warn!(target: "ias::sandbox", "容器不可用,降级本地执行");
}
// 降级:本地 spawn(审批已在入口完成,安全边界不受影响)
build_local_cmd(spec, binary_path)
.stdin(std::process::Stdio::piped())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.kill_on_drop(true)
.spawn()
}
+1
View File
@@ -6,4 +6,5 @@ pub mod approval;
pub mod executor; pub mod executor;
pub mod parser; pub mod parser;
pub mod registry; pub mod registry;
pub mod sandbox;
pub mod spec; pub mod spec;
+26 -17
View File
@@ -8,7 +8,7 @@ use std::sync::Arc;
use crate::context::MemoryStore; use crate::context::MemoryStore;
use crate::tools::approval::ApprovalManager; use crate::tools::approval::ApprovalManager;
use crate::tools::executor; use crate::tools::executor::{self, ToolErrorKind, ToolRunOutcome};
use crate::tools::parser; use crate::tools::parser;
use crate::tools::spec::ToolSpec; use crate::tools::spec::ToolSpec;
use serde_json::Value; use serde_json::Value;
@@ -29,7 +29,7 @@ pub struct ParamSummary {
} }
/// 解析工具目录。相对路径基于二进制所在目录。 /// 解析工具目录。相对路径基于二进制所在目录。
fn resolve_dir(dir: &str) -> PathBuf { pub(crate) fn resolve_dir(dir: &str) -> PathBuf {
let path = Path::new(dir); let path = Path::new(dir);
if path.is_absolute() { if path.is_absolute() {
return path.to_path_buf(); return path.to_path_buf();
@@ -89,15 +89,22 @@ impl Registry {
} }
for mut spec in result.specs { for mut spec in result.specs {
// 解析二进制路径 // 优先使用统一产物目录 tools/bin/<tool>,不存在则 fallback 到工具子目录
let binary = match &spec.path { let binary = {
Some(p) => { let unified = root.join("bin").join(&spec.tool);
let p = p.trim_start_matches('/'); if unified.exists() {
path.join(p) unified
} } else {
None => { match &spec.path {
// 默认: {tool_dir}/target/release/{tool} Some(p) => {
path.join("target/release").join(&spec.tool) let p = p.trim_start_matches('/');
path.join(p)
}
None => {
// 默认: {tool_dir}/target/release/{tool}
path.join("target/release").join(&spec.tool)
}
}
} }
}; };
@@ -153,16 +160,16 @@ impl Registry {
+ Send + Send
+ Sync + Sync
), ),
) -> String { ) -> ToolRunOutcome {
let spec = match self.get(name) { let spec = match self.get(name) {
Some(s) => s, Some(s) => s,
None => { None => {
tracing::warn!(target:"ias::registry","未知: {name}"); tracing::warn!(target:"ias::registry","未知: {name}");
return format!("未知工具: {name}"); return ToolRunOutcome::err(ToolErrorKind::UnknownTool, format!("未知工具: {name}"));
} }
}; };
if let Some(err) = validate_required_params(spec, params) { if let Some(err) = validate_required_params(spec, params) {
return err; return ToolRunOutcome::err(ToolErrorKind::MissingParams, err);
} }
executor::execute_loop( executor::execute_loop(
spec, spec,
@@ -215,20 +222,20 @@ impl RegistryManager {
+ Send + Send
+ Sync + Sync
), ),
) -> String { ) -> ToolRunOutcome {
let spec = { let spec = {
let registry = self.registry.read().await; let registry = self.registry.read().await;
match registry.get(name) { match registry.get(name) {
Some(spec) => spec.clone(), Some(spec) => spec.clone(),
None => { None => {
tracing::warn!(target:"ias::registry","未知: {name}"); tracing::warn!(target:"ias::registry","未知: {name}");
return format!("未知工具: {name}"); return ToolRunOutcome::err(ToolErrorKind::UnknownTool, format!("未知工具: {name}"));
} }
} }
}; };
if let Some(err) = validate_required_params(&spec, params) { if let Some(err) = validate_required_params(&spec, params) {
return err; return ToolRunOutcome::err(ToolErrorKind::MissingParams, err);
} }
executor::execute_loop( executor::execute_loop(
@@ -340,6 +347,7 @@ path: /target/release/mytool
required: true, required: true,
desc: "城市".into(), desc: "城市".into(),
}], }],
sandbox: Default::default(),
}; };
let err = validate_required_params(&spec, &serde_json::json!({})).unwrap(); let err = validate_required_params(&spec, &serde_json::json!({})).unwrap();
@@ -363,6 +371,7 @@ path: /target/release/mytool
required: true, required: true,
desc: "城市".into(), desc: "城市".into(),
}], }],
sandbox: Default::default(),
}; };
assert!( assert!(
+216
View File
@@ -0,0 +1,216 @@
//! # 沙箱容器管理
//!
//! 长期容器模式:daemon 启动时创建 `ias-sandbox` 容器(无网络、只读根、cap-drop ALL),
//! 统一产物目录 `tools/bin/` 只读挂载到容器 `/tools`。工具调用通过 `docker exec`
//! 在容器内执行,相比每次 `docker run` 省去容器创建/销毁开销。
//!
//! ## 超时与清理
//! `docker exec` 内用 `timeout -k 5 <secs>` 包裹工具进程:超时发 SIGTERM,
//! 5 秒后仍不退出发 SIGKILL,确保容器内不累积孤儿进程。executor 的 tokio timeout
//! 作为外层保险,kill `docker exec` 客户端释放宿主资源。
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};
use tokio::process::Command;
use crate::tools::spec::ToolSpec;
/// 长期沙箱容器名
pub const CONTAINER_NAME: &str = "ias-sandbox";
/// 容器内产物挂载点
const MOUNT_TARGET: &str = "/tools";
/// 降级标志:ensure_container 失败或容器崩溃后置 true,
/// executor 据此回退本地执行,避免不受信任工具持续 SpawnFailed
static DEGRADED: AtomicBool = AtomicBool::new(false);
/// 全局沙箱开关:`IAS_DOCKER_SANDBOX=1` 启用
pub fn enabled() -> bool {
std::env::var("IAS_DOCKER_SANDBOX")
.ok()
.as_deref()
== Some("1")
}
/// 当前是否处于降级状态(ensure 失败或容器崩溃)
pub fn degraded() -> bool {
DEGRADED.load(Ordering::SeqCst)
}
/// 执行前检查:沙箱启用且未降级时才走 docker exec
pub fn should_use_docker(profile: &str) -> bool {
enabled() && profile != "local" && !degraded()
}
/// 确保 ias-sandbox 长期容器在运行。daemon 启动时调用。
///
/// - 容器已运行:直接复用
/// - 容器不存在/已停止:清理后重新创建
///
/// 失败时设置降级标志,executor 据此回退本地执行。
pub async fn ensure_container(dir: &str) -> Result<(), String> {
match ensure_container_inner(dir).await {
Ok(()) => {
DEGRADED.store(false, Ordering::SeqCst);
Ok(())
}
Err(e) => {
DEGRADED.store(true, Ordering::SeqCst);
tracing::warn!(target: "ias::sandbox", "沙箱降级为本地执行: {e}");
Err(e)
}
}
}
async fn ensure_container_inner(dir: &str) -> Result<(), String> {
let root = locate_tools_root(dir)?;
let bin_dir = root.join("bin");
std::fs::create_dir_all(&bin_dir).map_err(|e| format!("创建产物目录失败: {e}"))?;
// 检查是否已在运行
let running = Command::new("docker")
.args(["inspect", "-f", "{{.State.Running}}", CONTAINER_NAME])
.output()
.await
.map(|o| String::from_utf8_lossy(&o.stdout).trim() == "true")
.unwrap_or(false);
if running {
tracing::info!(target: "ias::sandbox", "沙箱容器已运行,复用");
return Ok(());
}
// 清理旧的已停止容器
let _ = Command::new("docker")
.args(["rm", "-f", CONTAINER_NAME])
.output()
.await;
// 启动长期容器
let output = Command::new("docker")
.arg("run")
.arg("-d")
.arg("--rm")
.arg("--name")
.arg(CONTAINER_NAME)
.arg("--network")
.arg("none")
.arg("--cap-drop")
.arg("ALL")
.arg("--security-opt")
.arg("no-new-privileges")
.arg("--read-only")
.arg("--memory")
.arg("512m")
.arg("--cpus")
.arg("1")
.arg("--pids-limit")
.arg("128")
.arg("--mount")
.arg(format!(
"type=bind,source={},target={MOUNT_TARGET},readonly",
bin_dir.display()
))
.arg("--mount")
.arg("type=tmpfs,target=/tmp")
.arg("ias-stdio-runner")
.args(["sleep", "infinity"])
.output()
.await
.map_err(|e| format!("启动沙箱容器失败: {e}"))?;
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
return Err(format!("沙箱容器启动失败: {}", stderr.trim()));
}
tracing::info!(
target: "ias::sandbox",
"沙箱容器已启动 ({} → {})",
bin_dir.display(),
MOUNT_TARGET
);
Ok(())
}
/// exec 失败时调用:检测容器是否崩溃,崩溃则尝试重建。
/// 重建成功返回 true(调用方可重试 exec);否则置降级标志返回 false。
pub async fn recover_if_crashed(dir: &str) -> bool {
let running = Command::new("docker")
.args(["inspect", "-f", "{{.State.Running}}", CONTAINER_NAME])
.output()
.await
.map(|o| String::from_utf8_lossy(&o.stdout).trim() == "true")
.unwrap_or(false);
if running {
// 容器在运行,exec 失败是其他原因(如工具二进制不存在),不降级
return false;
}
tracing::warn!(target: "ias::sandbox", "沙箱容器已崩溃,尝试重建");
ensure_container(dir).await.is_ok()
}
/// 执行前调用:沙箱启用且未降级时,确保容器在运行。
/// 容器不在运行则重建;重建失败置降级标志并返回 false(调用方降级本地执行)。
pub async fn ensure_ready_or_degrade(dir: &str) -> bool {
if !enabled() || degraded() {
return false;
}
let running = Command::new("docker")
.args(["inspect", "-f", "{{.State.Running}}", CONTAINER_NAME])
.output()
.await
.map(|o| String::from_utf8_lossy(&o.stdout).trim() == "true")
.unwrap_or(false);
if running {
return true;
}
// 容器不在运行,尝试重建
if ensure_container(dir).await.is_ok() {
true
} else {
tracing::warn!(target: "ias::sandbox", "容器不可用,降级本地执行");
false
}
}
/// 构建 `docker exec` 命令,在长期容器内执行工具。
///
/// 容器内路径 = MOUNT_TARGET + 二进制文件名。
/// 用 `timeout -k 5` 包裹:超时 SIGTERM5 秒后 SIGKILL,避免孤儿进程。
/// env 按 spec 白名单注入(`-e` 必须在容器名前)。
pub fn build_exec_cmd(spec: &ToolSpec) -> Command {
// 直接用 spec.tool 作为容器内文件名,语义清晰,不依赖 spec.path 格式
let container_path = format!("{MOUNT_TARGET}/{}", spec.tool);
let mut cmd = Command::new("docker");
cmd.arg("exec").arg("-i");
// env 白名单注入(OPTIONS,必须在容器名前)
for key in &spec.env {
if let Ok(v) = std::env::var(key) {
cmd.arg("-e").arg(format!("{key}={v}"));
}
}
cmd.arg(CONTAINER_NAME);
// timeout -k 5 <secs>:超时 SIGTERM5s 后 SIGKILL
cmd.arg("timeout").arg("-k").arg("5").arg(spec.timeout_secs.to_string());
cmd.arg(container_path);
if let Some(c) = &spec.command {
cmd.arg("--command").arg(c);
}
cmd
}
/// 复用 registry 的路径解析逻辑定位 tools 根目录并转绝对路径
fn locate_tools_root(dir: &str) -> Result<PathBuf, String> {
let resolved = crate::tools::registry::resolve_dir(dir);
resolved
.canonicalize()
.map_err(|e| format!("无法定位 tools 目录 {}: {e}", resolved.display()))
}
+46
View File
@@ -13,6 +13,48 @@ pub struct ParamSpec {
pub desc: String, pub desc: String,
} }
/// 沙箱配置(Docker 一次性容器隔离)
///
/// ```yaml
/// sandbox:
/// profile: stdio # local | stdio | builder
/// network: none # none | default
/// memory: 256m
/// cpus: "0.5"
/// pids_limit: 64
/// read_only: true
/// ```
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(default)]
pub struct SandboxSpec {
/// 沙箱画像:`local` 本地执行(默认,不启用容器);`stdio` 无网络只读容器
/// (供 tool_manager 产出的不受信任工具);`builder` 构建容器(供 tool_manager
pub profile: String,
/// 网络:`none`(默认)或 `default`
pub network: String,
/// 内存限制,如 `256m`
pub memory: String,
/// CPU 限制,如 `0.5`
pub cpus: String,
/// PID 数量上限
pub pids_limit: u32,
/// 只读根文件系统
pub read_only: bool,
}
impl Default for SandboxSpec {
fn default() -> Self {
Self {
profile: "local".into(),
network: "none".into(),
memory: "256m".into(),
cpus: "0.5".into(),
pids_limit: 64,
read_only: true,
}
}
}
/// 工具规范(`*.tool.yaml` 的完整映射) /// 工具规范(`*.tool.yaml` 的完整映射)
/// ///
/// ```yaml /// ```yaml
@@ -56,6 +98,10 @@ pub struct ToolSpec {
#[serde(default)] #[serde(default)]
pub params: Vec<ParamSpec>, pub params: Vec<ParamSpec>,
/// 沙箱配置(Docker 隔离);默认 profile=local 不启用容器
#[serde(default)]
pub sandbox: SandboxSpec,
} }
fn default_risk_level() -> String { fn default_risk_level() -> String {
+1
View File
@@ -15,3 +15,4 @@ params:
- name: city - name: city
required: false required: false
desc: 城市名称,可选。限定地理编码的城市范围 desc: 城市名称,可选。限定地理编码的城市范围
+1
View File
@@ -12,3 +12,4 @@ params:
- name: center - name: center
required: false required: false
desc: 地图中心点经纬度坐标,格式"经度,纬度",默认天安门坐标 desc: 地图中心点经纬度坐标,格式"经度,纬度",默认天安门坐标
@@ -18,3 +18,4 @@ params:
- name: radius - name: radius
required: false required: false
desc: 搜索半径(米),如 500、1000、3000。默认 1000 米 desc: 搜索半径(米),如 500、1000、3000。默认 1000 米
@@ -12,3 +12,4 @@ params:
- name: location - name: location
required: true required: true
desc: 经纬度坐标,格式为"经度,纬度",如"116.397,39.908" desc: 经纬度坐标,格式为"经度,纬度",如"116.397,39.908"
@@ -18,3 +18,4 @@ params:
- name: type - name: type
required: false required: false
desc: 出行方式:driving(驾车)、walking(步行)、riding(骑行)、transit(公交)。默认 driving desc: 出行方式:driving(驾车)、walking(步行)、riding(骑行)、transit(公交)。默认 driving
@@ -15,3 +15,4 @@ params:
- name: interests - name: interests
required: false required: false
desc: 感兴趣的关键词,逗号分隔,如"景点,美食,酒店"。默认"景点,美食" desc: 感兴趣的关键词,逗号分隔,如"景点,美食,酒店"。默认"景点,美食"
+5 -4
View File
@@ -1,17 +1,18 @@
#!/bin/bash #!/bin/bash
# 构建所有工具 # 构建所有工具,产物统一收集到 bin/ 供长期沙箱容器挂载
set -e set -e
cd "$(dirname "$0")" cd "$(dirname "$0")"
TOOLS=(datetime weather amap web_search fetch_page memories tool_manager) TOOLS=(datetime weather amap web_search fetch_page memories tool_manager)
mkdir -p bin
echo "=== 构建工具 ===" echo "=== 构建工具 ==="
for tool in "${TOOLS[@]}"; do for tool in "${TOOLS[@]}"; do
echo " $tool..." echo " $tool..."
(cd "$tool" && cargo build --release --quiet) (cd "$tool" && cargo build --release --quiet)
cp "$tool/target/release/$tool" "bin/$tool"
done done
echo "=== 完成 ===" echo "=== 完成 ==="
for tool in "${TOOLS[@]}"; do ls -lh bin/*
ls -lh "$tool/target/release/$tool" 2>/dev/null
done
@@ -6,3 +6,4 @@ path: /target/release/datetime
risk_level: low risk_level: low
timeout_secs: 3 timeout_secs: 3
params: [] params: []
@@ -9,3 +9,4 @@ params:
- name: url - name: url
required: true required: true
desc: 网页URL desc: 网页URL
@@ -12,3 +12,4 @@ params:
- name: memory_id - name: memory_id
required: true required: true
desc: 记忆ID desc: 记忆ID
@@ -12,3 +12,4 @@ params:
- name: memory_id - name: memory_id
required: true required: true
desc: 记忆的数字ID,从 read_memories 返回结果中获取 desc: 记忆的数字ID,从 read_memories 返回结果中获取
@@ -9,3 +9,4 @@ timeout_secs: 3
env: env:
- IAS_MEMORY_DIR - IAS_MEMORY_DIR
params: params:
@@ -15,3 +15,4 @@ params:
- name: content - name: content
required: true required: true
desc: 更新后的新内容,覆盖原记忆 desc: 更新后的新内容,覆盖原记忆
@@ -12,3 +12,4 @@ params:
- name: content - name: content
required: true required: true
desc: 要记录的记忆内容,用自然语言描述。如"用户住在合肥包河区联投新安里A区" desc: 要记录的记忆内容,用自然语言描述。如"用户住在合肥包河区联投新安里A区"
+105
View File
@@ -0,0 +1,105 @@
# This file is automatically @generated by Cargo.
# It is not intended for manual editing.
version = 4
[[package]]
name = "itoa"
version = "1.0.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682"
[[package]]
name = "memchr"
version = "2.8.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4"
[[package]]
name = "proc-macro2"
version = "1.0.106"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934"
dependencies = [
"unicode-ident",
]
[[package]]
name = "quote"
version = "1.0.46"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368"
dependencies = [
"proc-macro2",
]
[[package]]
name = "serde"
version = "1.0.228"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e"
dependencies = [
"serde_core",
]
[[package]]
name = "serde_core"
version = "1.0.228"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad"
dependencies = [
"serde_derive",
]
[[package]]
name = "serde_derive"
version = "1.0.228"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "serde_json"
version = "1.0.150"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9"
dependencies = [
"itoa",
"memchr",
"serde",
"serde_core",
"zmij",
]
[[package]]
name = "syn"
version = "2.0.118"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422"
dependencies = [
"proc-macro2",
"quote",
"unicode-ident",
]
[[package]]
name = "tool_inspector"
version = "0.1.0"
dependencies = [
"serde_json",
]
[[package]]
name = "unicode-ident"
version = "1.0.24"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
[[package]]
name = "zmij"
version = "1.0.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa"
+11
View File
@@ -0,0 +1,11 @@
[package]
name = "tool_inspector"
version = "0.1.0"
edition = "2024"
[[bin]]
name = "tool_inspector"
path = "src/main.rs"
[dependencies]
serde_json = "1.0"
@@ -0,0 +1,17 @@
name: tool_inspector
desc: 只读查看本地工具的源码、Cargo.toml 和 spec 配置。在用 tool_manager 修改某工具前,先用本工具查看其当前实现,避免盲目覆盖丢失逻辑。低风险,无需审批
type: stdio
tool: tool_inspector
path: /target/release/tool_inspector
risk_level: low
timeout_secs: 10
params:
- name: action
required: true
desc: 操作类型:list(列出所有工具及文件构成)、read_tool(读取某工具的全部 Cargo.toml/specs/src 源码)、read_file(读取工具目录下单个文件)
- name: tool_name
required: false
desc: 工具名称(read_tool/read_file 必填),只允许字母、数字、下划线和短横线
- name: file
required: false
desc: read_file 时的相对路径,如 src/main.rs、Cargo.toml、specs/weather.tool.yaml;必须是相对路径,禁止 .. 和 target 目录
+270
View File
@@ -0,0 +1,270 @@
//! tool_inspector — 只读查看本地工具的源码与配置
//!
//! 用途:在用 tool_manager 修改工具前,先用本工具查看当前实现,避免盲目覆盖丢失逻辑。
//!
//! 安全约束:
//! - 纯只读:不写文件、不执行命令、不联网
//! - tool_name 与相对路径严格校验,禁止路径穿越
//! - 跳过 target/ 构建产物目录
//! - 低风险(risk_level: low),无需用户审批
use serde_json::{Value, json};
use std::fs;
use std::path::{Component, Path, PathBuf};
fn main() {
let input: Value = match serde_json::from_reader(std::io::stdin()) {
Ok(v) => v,
Err(e) => return result(&format!("读取输入失败: {e}")),
};
let params = &input["params"];
let action = params["action"].as_str().unwrap_or("");
let out = match run(action, params) {
Ok(s) => s,
Err(e) => format!("查看失败: {e}"),
};
result(&out);
}
fn run(action: &str, params: &Value) -> Result<String, String> {
let tools_root = locate_tools_root()?;
match action {
"list" => list_tools(&tools_root),
"read_tool" => {
let name = required_str(params, "tool_name")?;
validate_tool_name(name)?;
read_tool(&tools_root, name)
}
"read_file" => {
let name = required_str(params, "tool_name")?;
validate_tool_name(name)?;
let file = required_str(params, "file")?;
read_file(&tools_root, name, file)
}
_ => Err("action 必须是 list、read_tool 或 read_file".to_string()),
}
}
/// 列出 tools 目录下所有工具及其文件构成
fn list_tools(tools_root: &Path) -> Result<String, String> {
let mut names: Vec<String> = Vec::new();
for entry in fs::read_dir(tools_root).map_err(|e| e.to_string())? {
let entry = entry.map_err(|e| e.to_string())?;
let path = entry.path();
if !path.is_dir() {
continue;
}
let name = entry.file_name().to_string_lossy().to_string();
// 跳过产物目录与隐藏目录
if name == "bin" || name.starts_with('.') {
continue;
}
// 只有含 specs/src/Cargo.toml 之一的目录才视为工具
if path.join("specs").exists()
|| path.join("src").exists()
|| path.join("Cargo.toml").exists()
{
names.push(name);
}
}
names.sort();
if names.is_empty() {
return Ok("tools 目录下没有工具".into());
}
let mut out = format!("本地工具列表({} 个):\n", names.len());
for n in &names {
let dir = tools_root.join(n);
let mark = |has: bool| if has { "" } else { "" };
out.push_str(&format!(
"{n} (src:{} specs:{} Cargo.toml:{})\n",
mark(dir.join("src").exists()),
mark(dir.join("specs").exists()),
mark(dir.join("Cargo.toml").exists()),
));
}
out.push_str("\n用 read_tool 查看某工具的全部源码与配置,或用 read_file 读取单个文件。");
Ok(out)
}
/// 一次性读取某工具的 Cargo.toml + 所有 specs + src 下所有 .rs 文件
fn read_tool(tools_root: &Path, name: &str) -> Result<String, String> {
let tool_dir = tools_root.join(name);
if !tool_dir.exists() {
return Err(format!("工具 {name} 不存在"));
}
let mut out = String::new();
if let Ok(c) = fs::read_to_string(tool_dir.join("Cargo.toml")) {
out.push_str(&format!("===== {name}/Cargo.toml =====\n{c}\n\n"));
}
// specs/*.tool.yaml(可能有多个文件),按文件名排序
let specs_dir = tool_dir.join("specs");
if specs_dir.exists() {
let mut specs: Vec<PathBuf> = fs::read_dir(&specs_dir)
.map_err(|e| e.to_string())?
.filter_map(|e| e.ok())
.map(|e| e.path())
.filter(|p| matches!(p.extension().and_then(|e| e.to_str()), Some("yaml" | "yml")))
.collect();
specs.sort();
for spec in specs {
let rel = spec.strip_prefix(&tool_dir).unwrap_or(&spec).display().to_string();
match fs::read_to_string(&spec) {
Ok(c) => out.push_str(&format!("===== {name}/{rel} =====\n{c}\n\n")),
Err(e) => out.push_str(&format!("===== {name}/{rel} =====\n(读取失败: {e}\n\n")),
}
}
}
// src 下所有 .rs 文件(递归,支持多模块工具)
let mut src_files: Vec<PathBuf> = Vec::new();
collect_rs_files(&tool_dir.join("src"), &mut src_files)?;
if src_files.is_empty() {
out.push_str(&format!("{name} 无 src/*.rs 文件)\n"));
} else {
src_files.sort();
for f in src_files {
let rel = f.strip_prefix(&tool_dir).unwrap_or(&f).display().to_string();
match fs::read_to_string(&f) {
Ok(c) => out.push_str(&format!("===== {name}/{rel} =====\n{c}\n\n")),
Err(e) => out.push_str(&format!("===== {name}/{rel} =====\n(读取失败: {e}\n\n")),
}
}
}
if out.is_empty() {
return Err(format!("工具 {name} 没有可读的源码或配置"));
}
Ok(out.trim_end().to_string())
}
/// 读取工具目录下指定的单个文件(严格防路径穿越)
fn read_file(tools_root: &Path, name: &str, file: &str) -> Result<String, String> {
let tool_dir = tools_root.join(name);
if !tool_dir.exists() {
return Err(format!("工具 {name} 不存在"));
}
let rel = sanitize_rel_path(file)?;
let target = tool_dir.join(&rel);
// 规范化后必须仍在 tool_dir 内(双重防穿越)
let canonical_tool = tool_dir.canonicalize().map_err(|e| e.to_string())?;
let canonical_target = target
.canonicalize()
.map_err(|_| format!("文件不存在: {file}"))?;
if !canonical_target.starts_with(&canonical_tool) {
return Err("路径越界,禁止访问工具目录之外的文件".into());
}
if canonical_target.is_dir() {
return Err(format!("{file} 是目录,请用 read_tool 查看或指定具体文件"));
}
let content = fs::read_to_string(&canonical_target)
.map_err(|e| format!("读取 {file} 失败: {e}"))?;
Ok(format!("===== {name}/{file} =====\n{content}"))
}
/// 递归收集 .rs 文件,跳过隐藏目录
fn collect_rs_files(dir: &Path, out: &mut Vec<PathBuf>) -> Result<(), String> {
if !dir.exists() {
return Ok(());
}
for entry in fs::read_dir(dir).map_err(|e| e.to_string())? {
let entry = entry.map_err(|e| e.to_string())?;
let path = entry.path();
if path.is_dir() {
if !entry.file_name().to_string_lossy().starts_with('.') {
collect_rs_files(&path, out)?;
}
} else if path.extension().and_then(|e| e.to_str()) == Some("rs") {
out.push(path);
}
}
Ok(())
}
/// 校验相对路径:禁止绝对路径、禁止 ..、禁止 target 构建产物
fn sanitize_rel_path(file: &str) -> Result<PathBuf, String> {
if file.is_empty() {
return Err("file 不能为空".into());
}
let p = PathBuf::from(file);
for comp in p.components() {
match comp {
Component::Normal(_) | Component::CurDir => {}
Component::ParentDir => return Err("file 禁止包含 ..".into()),
Component::RootDir | Component::Prefix(_) => {
return Err("file 必须是相对路径".into())
}
}
}
if p.starts_with("target") {
return Err("禁止读取 target 构建产物目录".into());
}
Ok(p)
}
fn validate_tool_name(name: &str) -> Result<(), String> {
if name.is_empty() {
return Err("缺少 tool_name".into());
}
if !name
.chars()
.all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
{
return Err("tool_name 只允许字母、数字、下划线和短横线".into());
}
Ok(())
}
fn required_str<'a>(params: &'a Value, key: &str) -> Result<&'a str, String> {
params[key]
.as_str()
.filter(|s| !s.is_empty())
.ok_or_else(|| format!("缺少参数 {key}"))
}
/// 定位 tools 根目录。
///
/// 产物可能位于两处之一:
/// - `tools/bin/tool_inspector`(统一产物目录,registry 优先使用)
/// - `tools/tool_inspector/target/release/tool_inspector`(构建产物)
///
/// 因此不依赖固定向上层数,而是从 exe 向上找第一个“含 bin/ 子目录且至少含一个
/// 工具子目录”的祖先,即为 tools 根。对两种位置均正确。
fn locate_tools_root() -> Result<PathBuf, String> {
let exe = std::env::current_exe().map_err(|e| e.to_string())?;
let mut cur = exe.parent();
while let Some(dir) = cur {
if dir.join("bin").is_dir() && has_tool_subdir(dir) {
return Ok(dir.to_path_buf());
}
cur = dir.parent();
}
Err("无法定位 tools 根目录".into())
}
/// 判断目录下是否存在至少一个工具子目录(含 specs/Cargo.toml/src 之一)
fn has_tool_subdir(dir: &Path) -> bool {
let Ok(entries) = fs::read_dir(dir) else {
return false;
};
for entry in entries.flatten() {
let p = entry.path();
if p.is_dir()
&& (p.join("specs").exists()
|| p.join("Cargo.toml").exists()
|| p.join("src").exists())
{
return true;
}
}
false
}
fn result(content: &str) {
println!("{}", json!({"type":"result","content":content}));
}
@@ -4,7 +4,11 @@ type: stdio
tool: tool_manager tool: tool_manager
path: /target/release/tool_manager path: /target/release/tool_manager
risk_level: high risk_level: high
timeout_secs: 120 timeout_secs: 300
env:
- HOME
- CARGO_HOME
- RUSTUP_HOME
params: params:
- name: action - name: action
required: true required: true
@@ -23,7 +27,7 @@ params:
desc: 参数列表数组,每项包含 name、required、desc desc: 参数列表数组,每项包含 name、required、desc
- name: risk_level - name: risk_level
required: false required: false
desc: 新工具风险级别,low 或 high,默认 low desc: 已弃用,tool_manager 创建/更新的工具一律强制 high risk,传入值将被忽略
- name: timeout_secs - name: timeout_secs
required: false required: false
desc: 新工具超时时间秒数,默认 30 desc: 新工具超时时间秒数,默认 30
+169 -30
View File
@@ -1,10 +1,22 @@
//! tool_manager — 受控创建、修改、构建本地工具 //! tool_manager — 受控创建、修改、构建本地工具
//!
//! 安全约束:
//! - 不允许操作保留工具名(含 `tool_manager` 自身),防止自改后门
//! - 由本工具创建/更新的工具一律 `risk_level: high`,忽略调用方传入值
//! - 所有调用追加写入 `tools/.tool_audit.log` 审计日志
//! - 构建使用 `timeout` 限制时长,失败时清理残留二进制,避免调用旧版本"假成功"
use serde_json::{Value, json}; use serde_json::{Value, json};
use std::fs; use std::fs;
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
use std::process::Command; use std::process::Command;
/// 保留工具名 —— 禁止 create/update/build
const RESERVED_NAMES: &[&str] = &["target", "specs", "src", ".", "..", "tool_manager"];
/// 构建超时(秒)。需小于 tool_manager 自身的 timeout_secs。
const BUILD_TIMEOUT_SECS: u64 = 180;
fn main() { fn main() {
let input: Value = match serde_json::from_reader(std::io::stdin()) { let input: Value = match serde_json::from_reader(std::io::stdin()) {
Ok(v) => v, Ok(v) => v,
@@ -12,13 +24,20 @@ fn main() {
}; };
let params = &input["params"]; let params = &input["params"];
let user_id = input["user_id"].as_str().unwrap_or("unknown");
let action = params["action"].as_str().unwrap_or(""); let action = params["action"].as_str().unwrap_or("");
let tool_name = params["tool_name"].as_str().unwrap_or(""); let tool_name = params["tool_name"].as_str().unwrap_or("");
let output = match run(action, tool_name, params) { let (output, success) = match run(action, tool_name, params) {
Ok(msg) => msg, Ok(msg) => (msg, true),
Err(e) => format!("工具管理失败: {e}"), Err(e) => (format!("工具管理失败: {e}"), false),
}; };
// 审计日志(best effort,不影响主流程)
if let Ok(tools_root) = locate_tools_root() {
audit(&tools_root, user_id, action, tool_name, success, &output);
}
result(&output); result(&output);
} }
@@ -38,7 +57,7 @@ fn run(action: &str, tool_name: &str, params: &Value) -> Result<String, String>
write_tool_project(&tool_dir, tool_name, params)?; write_tool_project(&tool_dir, tool_name, params)?;
build_tool(&tool_dir, tool_name)?; build_tool(&tool_dir, tool_name)?;
Ok(format!( Ok(format!(
"已创建并构建工具 {tool_name}。可通过 query_capabilities 查看并用 call_capability 调用。" "已创建并构建工具 {tool_name}risk_level 强制为 high。可通过 query_capabilities 查看并用 call_capability 调用。"
)) ))
} }
"update_tool" => { "update_tool" => {
@@ -48,7 +67,7 @@ fn run(action: &str, tool_name: &str, params: &Value) -> Result<String, String>
update_tool_project(&tool_dir, tool_name, params)?; update_tool_project(&tool_dir, tool_name, params)?;
build_tool(&tool_dir, tool_name)?; build_tool(&tool_dir, tool_name)?;
Ok(format!( Ok(format!(
"已修改并构建工具 {tool_name}。注册表重载后即可使用最新版本。" "已修改并构建工具 {tool_name}risk_level 强制为 high。注册表重载后即可使用最新版本。"
)) ))
} }
"build_tool" => { "build_tool" => {
@@ -83,6 +102,7 @@ fn update_tool_project(tool_dir: &Path, tool_name: &str, params: &Value) -> Resu
fs::write(tool_dir.join("src/main.rs"), code).map_err(|e| e.to_string())?; fs::write(tool_dir.join("src/main.rs"), code).map_err(|e| e.to_string())?;
} }
// spec 一旦由 tool_manager 管理,risk_level 恒为 high
let should_update_spec = params.get("description").is_some() let should_update_spec = params.get("description").is_some()
|| params.get("params").is_some() || params.get("params").is_some()
|| params.get("risk_level").is_some() || params.get("risk_level").is_some()
@@ -101,36 +121,109 @@ fn update_tool_project(tool_dir: &Path, tool_name: &str, params: &Value) -> Resu
} }
fn build_tool(tool_dir: &Path, tool_name: &str) -> Result<(), String> { fn build_tool(tool_dir: &Path, tool_name: &str) -> Result<(), String> {
let output = Command::new("cargo") let binary = tool_dir.join("target/release").join(tool_name);
.arg("build") let use_builder = std::env::var("IAS_DOCKER_BUILDER")
.arg("--release") .ok()
.current_dir(tool_dir) .as_deref()
.output() == Some("1");
.map_err(|e| format!("启动 cargo 失败: {e}"))?; let (uid, gid) = current_uid_gid();
// 启用 builder 容器时在隔离环境内离线编译,挂载 tool_dir 到 /work,产物写回宿主;
// 否则本地 cargo(受 timeout 限制)
let output = if use_builder {
Command::new("docker")
.arg("run").arg("--rm")
.arg("--network").arg("none")
.arg("--memory").arg("1g")
.arg("--cpus").arg("1")
.arg("--pids-limit").arg("128")
.arg("--cap-drop").arg("ALL")
.arg("--security-opt").arg("no-new-privileges")
.arg("--user").arg(format!("{}:{}", uid, gid))
.arg("-v").arg(format!("{}:/work", tool_dir.display()))
.arg("-w").arg("/work")
.arg("-e").arg("CARGO_NET_OFFLINE=true")
.arg("-e").arg("CARGO_HOME=/tmp/cargo")
.arg("ias-tool-builder")
.arg("timeout").arg(BUILD_TIMEOUT_SECS.to_string())
.arg("cargo").arg("build").arg("--release").arg("--offline")
.output()
} else {
Command::new("timeout")
.arg(BUILD_TIMEOUT_SECS.to_string())
.arg("cargo")
.arg("build")
.arg("--release")
.current_dir(tool_dir)
.output()
}
.map_err(|e| format!("启动构建失败: {e}"))?;
if !output.status.success() { if !output.status.success() {
// 构建失败:删除可能残留的旧二进制,避免 reload 后仍调用旧版本造成"假成功"
let _ = fs::remove_file(&binary);
let code = output.status.code().unwrap_or(-1);
let stderr = String::from_utf8_lossy(&output.stderr); let stderr = String::from_utf8_lossy(&output.stderr);
return Err(format!("cargo build 失败: {}", stderr.trim())); let reason = if code == 124 {
format!("构建超时({BUILD_TIMEOUT_SECS}s")
} else {
stderr.trim().to_string()
};
return Err(format!("cargo build 失败(退出码 {code}): {reason}"));
} }
let binary = tool_dir.join("target/release").join(tool_name);
if !binary.exists() { if !binary.exists() {
return Err(format!("构建后未找到二进制: {}", binary.display())); return Err(format!("构建后未找到二进制: {}", binary.display()));
} }
// 复制到统一产物目录 tools/bin/,供长期沙箱容器挂载使用
let tools_root = tool_dir
.parent()
.ok_or_else(|| "无法定位 tools 根目录".to_string())?;
let bin_dir = tools_root.join("bin");
fs::create_dir_all(&bin_dir).map_err(|e| format!("创建产物目录失败: {e}"))?;
fs::copy(&binary, bin_dir.join(tool_name))
.map_err(|e| format!("复制产物到 bin 失败: {e}"))?;
Ok(()) Ok(())
} }
/// 定位 tools 根目录。
///
/// 产物可能位于两处之一:
/// - `tools/bin/tool_manager`(统一产物目录,registry 优先使用)
/// - `tools/tool_manager/target/release/tool_manager`(构建产物)
///
/// 因此不依赖固定向上层数,而是从 exe 向上找第一个“含 bin/ 子目录且至少含一个
/// 工具子目录”的祖先,即为 tools 根。对两种位置均正确。
fn locate_tools_root() -> Result<PathBuf, String> { fn locate_tools_root() -> Result<PathBuf, String> {
let exe = std::env::current_exe().map_err(|e| e.to_string())?; let exe = std::env::current_exe().map_err(|e| e.to_string())?;
let tool_dir = exe let mut cur = exe.parent();
.parent() while let Some(dir) = cur {
.and_then(Path::parent) if dir.join("bin").is_dir() && has_tool_subdir(dir) {
.and_then(Path::parent) return Ok(dir.to_path_buf());
.ok_or_else(|| "无法定位 tool_manager 目录".to_string())?; }
tool_dir cur = dir.parent();
.parent() }
.map(Path::to_path_buf) Err("无法定位 tools 根目录".into())
.ok_or_else(|| "无法定位 tools 根目录".to_string()) }
/// 判断目录下是否存在至少一个工具子目录(含 specs/Cargo.toml/src 之一)
fn has_tool_subdir(dir: &Path) -> bool {
let Ok(entries) = fs::read_dir(dir) else {
return false;
};
for entry in entries.flatten() {
let p = entry.path();
if p.is_dir()
&& (p.join("specs").exists()
|| p.join("Cargo.toml").exists()
|| p.join("src").exists())
{
return true;
}
}
false
} }
fn cargo_toml(tool_name: &str) -> String { fn cargo_toml(tool_name: &str) -> String {
@@ -171,14 +264,13 @@ fn main() {{
} }
fn spec_yaml(tool_name: &str, params: &Value) -> String { fn spec_yaml(tool_name: &str, params: &Value) -> String {
// 安全约束:tool_manager 产出的工具一律高风险,忽略调用方传入的 risk_level
let risk_level = "high";
let timeout_secs = params["timeout_secs"].as_u64().unwrap_or(30).clamp(1, 300);
let desc = params["description"] let desc = params["description"]
.as_str() .as_str()
.unwrap_or("由 tool_manager 创建的本地工具"); .unwrap_or("由 tool_manager 创建的本地工具(自动标记 high risk");
let risk_level = match params["risk_level"].as_str() {
Some("high") => "high",
_ => "low",
};
let timeout_secs = params["timeout_secs"].as_u64().unwrap_or(30).clamp(1, 300);
let mut out = format!( let mut out = format!(
"name: {tool_name}\n\ "name: {tool_name}\n\
@@ -211,6 +303,10 @@ fn spec_yaml(tool_name: &str, params: &Value) -> String {
out.push_str("params: []\n"); out.push_str("params: []\n");
} }
// 安全约束:tool_manager 产出的工具不受信任(含 LLM 提供的 rust_code),
// 默认进 stdio 沙箱容器(无网络、只读根),防止数据外泄/越权文件读
out.push_str("sandbox:\n profile: stdio\n");
out out
} }
@@ -218,8 +314,10 @@ fn validate_tool_name(name: &str) -> Result<(), String> {
if name.is_empty() { if name.is_empty() {
return Err("缺少 tool_name".to_string()); return Err("缺少 tool_name".to_string());
} }
if matches!(name, "target" | "specs" | "src" | "." | "..") { if RESERVED_NAMES.contains(&name) {
return Err("tool_name 使用了保留名称".to_string()); return Err(format!(
"tool_name '{name}' 是保留名称,禁止 create/update/build"
));
} }
if !name if !name
.chars() .chars()
@@ -253,6 +351,47 @@ fn truthy(value: Option<&Value>) -> bool {
} }
} }
/// 获取当前进程的 uid:gid,用于 builder 容器 --user,使容器内进程能写 bind mount 的宿主目录
fn current_uid_gid() -> (String, String) {
let uid = String::from_utf8_lossy(
&Command::new("id").arg("-u").output().map(|o| o.stdout).unwrap_or_default(),
)
.trim()
.to_string();
let gid = String::from_utf8_lossy(
&Command::new("id").arg("-g").output().map(|o| o.stdout).unwrap_or_default(),
)
.trim()
.to_string();
(uid, gid)
}
/// 追加审计日志到 tools 根目录下的 .tool_audit.log
fn audit(
tools_root: &Path,
user_id: &str,
action: &str,
tool_name: &str,
success: bool,
detail: &str,
) {
let ts = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0);
// detail 可能很长,截断到 200 字节避免日志膨胀
let detail: String = detail.chars().take(200).collect();
let line = format!(
"{ts}\t{user_id}\t{action}\t{tool_name}\t{}\t{detail}\n",
if success { "ok" } else { "err" }
);
let path = tools_root.join(".tool_audit.log");
if let Ok(mut f) = fs::OpenOptions::new().create(true).append(true).open(&path) {
use std::io::Write;
let _ = f.write_all(line.as_bytes());
}
}
fn result(content: &str) { fn result(content: &str) {
println!("{}", json!({"type":"result","content":content})); println!("{}", json!({"type":"result","content":content}));
} }
@@ -12,3 +12,4 @@ params:
- name: location - name: location
required: true required: true
desc: 城市名称,如"北京"、"上海"、"合肥"。支持地级市和县级市 desc: 城市名称,如"北京"、"上海"、"合肥"。支持地级市和县级市
@@ -12,3 +12,4 @@ params:
- name: location - name: location
required: true required: true
desc: 城市名称,如"北京"、"上海"、"合肥"。支持地级市和县级市 desc: 城市名称,如"北京"、"上海"、"合肥"。支持地级市和县级市
@@ -12,3 +12,4 @@ params:
- name: location - name: location
required: true required: true
desc: 城市名称,如"北京"、"上海"、"合肥"。支持地级市和县级市 desc: 城市名称,如"北京"、"上海"、"合肥"。支持地级市和县级市
@@ -13,3 +13,4 @@ params:
desc: 搜索关键词,如"Rust 2024 edition 新特性"、"今天比特币价格" desc: 搜索关键词,如"Rust 2024 edition 新特性"、"今天比特币价格"
- name: max_results - name: max_results
desc: 返回的最大搜索结果数量(1-20),默认5条 desc: 返回的最大搜索结果数量(1-20),默认5条