Compare commits
12 Commits
32c4fcf3a8
..
main
| Author | SHA1 | Date | |
|---|---|---|---|
| 3476d8cbcb | |||
| 132089c8eb | |||
| 0299329fb2 | |||
| a2f88b9b01 | |||
| 142efb070f | |||
| 092f3b78f2 | |||
| 4a5b21a69e | |||
| cacd4436ce | |||
| 20f7c3ba9c | |||
| 1e33f527e0 | |||
| 4f80e81b0b | |||
| bf231352fb |
@@ -1,11 +1,14 @@
|
||||
# Rust 构建产物
|
||||
/target
|
||||
tools/*/target/
|
||||
tools/bin/
|
||||
dist/
|
||||
|
||||
# 运行时数据(状态文件、缓存、聊天记录本地副本等)
|
||||
.data/
|
||||
.pi
|
||||
# tool_manager 审计日志
|
||||
tools/.tool_audit.log
|
||||
|
||||
# 环境配置(含密钥,不入库)
|
||||
.env
|
||||
|
||||
@@ -333,6 +333,33 @@ flowchart TB
|
||||
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. 文件依赖矩阵
|
||||
|
||||
@@ -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 []
|
||||
@@ -0,0 +1,13 @@
|
||||
# ias-tool-builder — tool_manager 编译 Rust 工具的隔离环境
|
||||
# 预缓存 serde_json(tool_manager 生成工具的唯一依赖),支持 --network none 离线编译。
|
||||
# CARGO_HOME=/tmp/cargo 并 chmod 777,使任意 uid(--user 运行)都能复用缓存。
|
||||
# 运行时由 tool_manager 挂载 tool_dir 到 /work,cargo 在 /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 []
|
||||
Executable
+16
@@ -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-builder(Rust 编译容器)==="
|
||||
docker build -t ias-tool-builder -f Dockerfile.tool-builder .
|
||||
|
||||
echo "=== 完成 ==="
|
||||
docker images | grep -E '^ias-(stdio|tool)'
|
||||
@@ -0,0 +1,186 @@
|
||||
# 两层安全边界设计
|
||||
|
||||
> 审批决定"能不能做",Docker 限制"做坏了能坏到哪里",审计回答"谁让系统新增或修改了什么能力"。
|
||||
|
||||
## 架构
|
||||
|
||||
```
|
||||
LLM 调用工具
|
||||
→ registry.dispatch(name, params, user_id, approved)
|
||||
→ load spec + validate params
|
||||
→ 【授权边界】if high_risk && !approved: 统一审批(不依赖消息类型)
|
||||
→ 【隔离边界】按 sandbox.profile 选择执行方式:
|
||||
- local : 本地 spawn(env_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 | 可写 /work(bind 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` 据此判断是否 reload(tool_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-sandbox,exec 完无孤儿进程 |
|
||||
| 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
@@ -260,13 +260,48 @@ pub async fn trigger_overflow_summary(
|
||||
/// 而是只保存到数据库,供 `read_summaries` 工具查询。
|
||||
///
|
||||
/// 生成摘要后会清空所有消息,checkpoint 重置为 0。
|
||||
/// 该摘要的 `created_at` 同时作为"会话边界"——之后加载历史时
|
||||
/// 不会包含此时间点之前的消息(见 `db::latest_session_boundary`)。
|
||||
pub async fn trigger_idle_summary(
|
||||
session: &Arc<Mutex<ChatSession>>,
|
||||
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;
|
||||
if s.messages.is_empty() {
|
||||
return;
|
||||
return None;
|
||||
}
|
||||
|
||||
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 {
|
||||
checkpoint: cp,
|
||||
text: summary_text.clone(),
|
||||
reason: super::types::SummaryReason::Timeout,
|
||||
reason,
|
||||
created_at: chrono::Utc::now(),
|
||||
});
|
||||
|
||||
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;
|
||||
|
||||
s.checkpoint = s.messages.len();
|
||||
@@ -296,6 +331,8 @@ pub async fn trigger_idle_summary(
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Some(summary_text)
|
||||
}
|
||||
|
||||
/// 生成摘要:优先使用 LLM,不可用时回退到简单截断
|
||||
@@ -431,4 +468,41 @@ mod tests {
|
||||
assert!(has_anchor);
|
||||
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());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -14,6 +14,7 @@
|
||||
//! 2. AI 回复 → `add_assistant()` → 存入 `messages`
|
||||
//! 3. Token 超 budget → `trigger_overflow_summary()` → 压缩 checkpoint 前的消息
|
||||
//! 4. 12h 空闲 → `trigger_idle_summary()` → 生成摘要,清空消息
|
||||
//! 5. `/clear` 指令 → `trigger_clear_summary()` → 生成摘要归档,会话过期
|
||||
|
||||
use crate::llm::types::Message;
|
||||
use chrono::{DateTime, Utc};
|
||||
@@ -30,15 +31,18 @@ pub struct HistoryEntry {
|
||||
|
||||
/// ## 摘要原因
|
||||
///
|
||||
/// 区分两种触发摘要的场景:
|
||||
/// 区分触发摘要的场景:
|
||||
/// * `Overflow` — 上下文 token 总数超过预算(token budget),触发压缩
|
||||
/// * `Timeout` — 距上条用户消息超过 12 小时空闲,触发会话总结
|
||||
/// * `Clear` — 用户发送 `/clear` 主动清空上下文,触发会话过期归档
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub enum SummaryReason {
|
||||
/// 上下文 token 超 budget 触发
|
||||
Overflow,
|
||||
/// 距上条用户消息超过 12 小时触发
|
||||
Timeout,
|
||||
/// 用户发送 `/clear` 主动清空上下文触发
|
||||
Clear,
|
||||
}
|
||||
|
||||
/// ## 一条摘要记录
|
||||
@@ -77,6 +81,7 @@ pub struct SummaryEntry {
|
||||
/// 2. AI 回复 → `add_assistant()` → 存入 `messages`
|
||||
/// 3. Token 超 budget → `trigger_overflow_summary()` → 压缩 checkpoint 前的消息
|
||||
/// 4. 12h 空闲 → `trigger_idle_summary()` → 生成摘要,清空消息
|
||||
/// 5. `/clear` 指令 → `trigger_clear_summary()` → 生成摘要归档,会话过期
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ChatSession {
|
||||
pub user_id: String,
|
||||
|
||||
+204
-38
@@ -39,6 +39,7 @@
|
||||
//! 采用统一的单进程 tokio task 架构。旧版 daemon+worker 分离模式(UDS IPC)已在 v2 中移除。
|
||||
|
||||
use crate::channel::ChannelId;
|
||||
use crate::context::builder;
|
||||
use crate::context::HistoryEntry;
|
||||
use crate::context::MemoryStore;
|
||||
use crate::db::Database;
|
||||
@@ -49,6 +50,7 @@ use crate::llm::{
|
||||
use crate::queue::{ConsumerChannels, EnqueueHandle, MessageKind, PipelineMessage, QueueRunner};
|
||||
|
||||
use crate::tools::approval::ApprovalManager;
|
||||
use crate::tools::executor::ToolRunOutcome;
|
||||
use crate::tools::registry::RegistryManager;
|
||||
use crate::wechat::client::WeChatClient;
|
||||
use std::collections::HashMap;
|
||||
@@ -58,6 +60,10 @@ use tokio::sync::Mutex;
|
||||
use tracing::{error, info, warn};
|
||||
use uuid::Uuid;
|
||||
|
||||
/// LLM Consumer 的会话缓存:`correlation_id → (Conversation, 创建时间, user_id)`。
|
||||
/// 额外记录 `user_id` 是为了在 `/clear` 时能按用户丢弃尚未完成的在途会话。
|
||||
type SessionMap = HashMap<Uuid, (Conversation, Instant, String)>;
|
||||
|
||||
/// ## Daemon 上下文 —— 消费者之间共享的全部状态
|
||||
///
|
||||
/// 这个结构体被所有三个消费者(LLM / Tool / Send)通过 `Arc` 共享访问。
|
||||
@@ -111,6 +117,15 @@ pub async fn run(database: &Arc<Database>, memory_store: &Arc<MemoryStore>) {
|
||||
database.pool().clone(),
|
||||
))));
|
||||
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 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
|
||||
.lock()
|
||||
.await
|
||||
.retain(|_, (conv, created)| {
|
||||
.retain(|_, (conv, created, _)| {
|
||||
let ttl_secs = if conv.has_pending_async() { 1800 } else { 300 };
|
||||
now.duration_since(*created).as_secs() < ttl_secs
|
||||
});
|
||||
@@ -260,9 +275,9 @@ fn spawn_consumers(
|
||||
enqueue: EnqueueHandle,
|
||||
) -> (
|
||||
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()));
|
||||
|
||||
let mut handles = Vec::new();
|
||||
@@ -302,7 +317,7 @@ async fn llm_consumer_loop(
|
||||
rx: &mut tokio::sync::mpsc::Receiver<PipelineMessage>,
|
||||
ctx: &DaemonCtx,
|
||||
enqueue: &EnqueueHandle,
|
||||
sessions: &Arc<Mutex<HashMap<Uuid, (Conversation, Instant)>>>,
|
||||
sessions: &Arc<Mutex<SessionMap>>,
|
||||
) {
|
||||
while let Some(msg) = rx.recv().await {
|
||||
let user_id = msg.channel.user_id.clone();
|
||||
@@ -314,6 +329,19 @@ async fn llm_consumer_loop(
|
||||
message_id: _,
|
||||
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;
|
||||
let history = load_history(&ctx.db, &user_id).await;
|
||||
@@ -418,7 +446,7 @@ async fn llm_consumer_loop(
|
||||
sessions
|
||||
.lock()
|
||||
.await
|
||||
.insert(correlation_id, (conv, Instant::now()));
|
||||
.insert(correlation_id, (conv, Instant::now(), user_id.clone()));
|
||||
|
||||
// 执行 LLM 对话
|
||||
run_llm_round(
|
||||
@@ -442,7 +470,7 @@ async fn llm_consumer_loop(
|
||||
// 从缓存恢复会话,注入工具结果并恢复 LLM 循环
|
||||
let should_resume = {
|
||||
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(
|
||||
&tool_call_id,
|
||||
&tool_name,
|
||||
@@ -492,13 +520,13 @@ async fn run_llm_round(
|
||||
correlation_id: Uuid,
|
||||
context_token: Option<String>,
|
||||
text: &str,
|
||||
sessions: &Arc<Mutex<HashMap<Uuid, (Conversation, Instant)>>>,
|
||||
sessions: &Arc<Mutex<SessionMap>>,
|
||||
enqueue: &EnqueueHandle,
|
||||
) {
|
||||
let result = {
|
||||
let map = sessions.lock().await;
|
||||
let conv = match map.get(&correlation_id) {
|
||||
Some((c, _)) => c,
|
||||
Some((c, _, _)) => c,
|
||||
None => {
|
||||
error!(target: "ias::err","[llm-session] conv不在缓存");
|
||||
return;
|
||||
@@ -535,13 +563,13 @@ async fn resume_llm_round(
|
||||
user_id: &str,
|
||||
correlation_id: Uuid,
|
||||
context_token: Option<String>,
|
||||
sessions: &Arc<Mutex<HashMap<Uuid, (Conversation, Instant)>>>,
|
||||
sessions: &Arc<Mutex<SessionMap>>,
|
||||
enqueue: &EnqueueHandle,
|
||||
) {
|
||||
let result = {
|
||||
let map = sessions.lock().await;
|
||||
let conv = match map.get(&correlation_id) {
|
||||
Some((c, _)) => c,
|
||||
Some((c, _, _)) => c,
|
||||
None => {
|
||||
error!(target: "ias::err","[llm-resume] conv不在缓存");
|
||||
return;
|
||||
@@ -579,7 +607,7 @@ async fn handle_chat_result(
|
||||
correlation_id: Uuid,
|
||||
context_token: Option<String>,
|
||||
result: ChatResult,
|
||||
sessions: &Arc<Mutex<HashMap<Uuid, (Conversation, Instant)>>>,
|
||||
sessions: &Arc<Mutex<SessionMap>>,
|
||||
enqueue: &EnqueueHandle,
|
||||
) {
|
||||
if result.has_pending_async {
|
||||
@@ -607,6 +635,111 @@ async fn handle_chat_result(
|
||||
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 ───
|
||||
|
||||
async fn tool_consumer_loop(
|
||||
@@ -692,7 +825,7 @@ async fn handle_tool_call(
|
||||
eq.enqueue(msg).await;
|
||||
})
|
||||
};
|
||||
let output = ctx
|
||||
let outcome = ctx
|
||||
.registry
|
||||
.dispatch(
|
||||
target,
|
||||
@@ -704,11 +837,17 @@ async fn handle_tool_call(
|
||||
&wechat_cb,
|
||||
)
|
||||
.await;
|
||||
if target == "tool_manager" && !is_tool_error(&output) {
|
||||
if target == "tool_manager" && !outcome.is_err() {
|
||||
let count = ctx.registry.reload().await;
|
||||
format!("{output}\n\n工具注册表已重载,当前可用工具数: {count}")
|
||||
format!(
|
||||
"{}\n\n工具注册表已重载,当前可用工具数: {count}",
|
||||
outcome.into_message()
|
||||
)
|
||||
} else {
|
||||
output
|
||||
if let ToolRunOutcome::Err { ref kind, .. } = outcome {
|
||||
tracing::warn!(target:"ias::err","[tool] {target} 执行失败: {kind:?}");
|
||||
}
|
||||
outcome.into_message()
|
||||
}
|
||||
} else {
|
||||
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;
|
||||
})
|
||||
};
|
||||
let output = ctx
|
||||
let outcome = ctx
|
||||
.registry
|
||||
.dispatch(
|
||||
&tool_name,
|
||||
@@ -746,11 +885,17 @@ async fn handle_tool_call(
|
||||
&wechat_cb,
|
||||
)
|
||||
.await;
|
||||
if tool_name == "tool_manager" && !is_tool_error(&output) {
|
||||
if tool_name == "tool_manager" && !outcome.is_err() {
|
||||
let count = ctx.registry.reload().await;
|
||||
format!("{output}\n\n工具注册表已重载,当前可用工具数: {count}")
|
||||
format!(
|
||||
"{}\n\n工具注册表已重载,当前可用工具数: {count}",
|
||||
outcome.into_message()
|
||||
)
|
||||
} 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;
|
||||
}
|
||||
|
||||
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) {
|
||||
let target = cp["name"].as_str().unwrap_or("");
|
||||
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
@@ -138,22 +138,31 @@ pub async fn insert_chat_record(
|
||||
}
|
||||
|
||||
/// 查询最近可进入 LLM 上下文的聊天记录(排除工具等待/审批等中间提示)。
|
||||
///
|
||||
/// 会自动应用"会话边界":最近一次 `/clear` 或 12h 空闲超时摘要之前的消息
|
||||
/// 已归档为摘要,不再进入新会话的上下文(见 [`latest_session_boundary`])。
|
||||
pub async fn list_recent_chat_records_for_context(
|
||||
pool: &PgPool,
|
||||
user_id: &str,
|
||||
limit: i64,
|
||||
) -> Result<Vec<ChatRecord>, String> {
|
||||
// 会话边界:最近一次"会话过期"摘要(/clear 或 12h 空闲超时)的时间点。
|
||||
// 边界之前的消息已归档为摘要,不再进入新会话的上下文。
|
||||
let boundary = latest_session_boundary(pool, user_id).await;
|
||||
|
||||
let records = sqlx::query_as::<_, ChatRecord>(
|
||||
r#"
|
||||
SELECT id, created_at, direction, user_id, account_id, text, source, context_token, message_id
|
||||
FROM chat_records
|
||||
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
|
||||
LIMIT $2
|
||||
LIMIT $3
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(boundary)
|
||||
.bind(limit)
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
@@ -162,6 +171,32 @@ pub async fn list_recent_chat_records_for_context(
|
||||
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 用量聚合统计
|
||||
|
||||
@@ -478,7 +478,10 @@ pub const DEFAULT_SYSTEM_PROMPT: &str = "\
|
||||
1. `query_capabilities` — 查询所有可用能力工具\n\
|
||||
2. `call_capability` — 调用能力工具,参数 {\"name\":\"工具名\",\"prompt\":{\"参数\":\"值\"}}\n\
|
||||
\n\
|
||||
流程:收到消息 → query_capabilities → call_capability → 回复\n\
|
||||
流程:收到消息 → query_capabilities → call_capability → 回复
|
||||
\
|
||||
\
|
||||
需要创建或修改本地工具时:先用 `tool_inspector`(action=read_tool)查看现有实现,再用 `tool_manager` 修改。`tool_manager` 是高风险操作,执行前会请求用户确认。\n\
|
||||
";
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
+7
-24
@@ -529,7 +529,7 @@ async fn cmd_tool_dynamic(name: &str, args: &[String]) {
|
||||
+ Send
|
||||
+ Sync
|
||||
) = &|_: &str, _: &str| Box::pin(async {});
|
||||
let result = registry
|
||||
let outcome = registry
|
||||
.dispatch(
|
||||
name,
|
||||
&serde_json::Value::Object(params),
|
||||
@@ -541,28 +541,11 @@ async fn cmd_tool_dynamic(name: &str, args: &[String]) {
|
||||
)
|
||||
.await;
|
||||
|
||||
// 检测是否执行出错(匹配 execute_loop 返回的所有错误前缀)
|
||||
let is_error = 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("工具 '");
|
||||
|
||||
if is_error {
|
||||
eprintln!("❌ {result}");
|
||||
std::process::exit(1);
|
||||
} else {
|
||||
println!("{result}");
|
||||
match outcome {
|
||||
crate::tools::executor::ToolRunOutcome::Ok(s) => println!("{s}"),
|
||||
crate::tools::executor::ToolRunOutcome::Err { message, .. } => {
|
||||
eprintln!("❌ {message}");
|
||||
std::process::exit(1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+113
-11
@@ -129,6 +129,33 @@ impl ApprovalManager {
|
||||
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
|
||||
pub async fn handle_reply(
|
||||
@@ -192,22 +219,32 @@ impl ApprovalManager {
|
||||
}
|
||||
|
||||
/// 清理过期审批(按时间判定,通知等待方 + 更新 DB)
|
||||
///
|
||||
/// 两阶段加锁:收集过期 code_hash → 更新 DB → 重新加锁按 code_hash 定位删除。
|
||||
/// 不依赖快照 index,避免并发回复导致 index 错位误删。
|
||||
pub async fn clean_expired(&self) {
|
||||
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;
|
||||
for (uid, entries) in map.iter() {
|
||||
for (i, e) in entries.iter().enumerate() {
|
||||
for entries in map.values() {
|
||||
for e in entries {
|
||||
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 {
|
||||
for (_uid, _i, code_hash) in &expired {
|
||||
for code_hash in &expired_hashes {
|
||||
let _ = sqlx::query(
|
||||
"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;
|
||||
for (uid, i, _) in expired.iter().rev() {
|
||||
if let Some(entries) = map.get_mut(uid)
|
||||
&& *i < entries.len()
|
||||
{
|
||||
let entry = entries.remove(*i);
|
||||
let _ = entry.sender.send(ApprovalDecision::Expired);
|
||||
for code_hash in &expired_hashes {
|
||||
for entries in map.values_mut() {
|
||||
if let Some(idx) = entries.iter().position(|e| &e.code_hash == code_hash) {
|
||||
let entry = entries.remove(idx);
|
||||
let _ = entry.sender.send(ApprovalDecision::Expired);
|
||||
}
|
||||
}
|
||||
}
|
||||
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_expired:tool_a 已被 remove,按 hash 查找应跳过;
|
||||
// tool_b 仍在且过期,应被正确标记 Expired(不会被误当成 tool_a 跳过)
|
||||
mgr.clean_expired().await;
|
||||
assert_eq!(rx_b.await.unwrap(), ApprovalDecision::Expired);
|
||||
}
|
||||
}
|
||||
|
||||
+193
-78
@@ -5,7 +5,11 @@
|
||||
//! - `{"type":"db",...}` — 请求数据库操作
|
||||
//! - `{"type":"result","content":"..."}` — 最终结果
|
||||
//!
|
||||
//! 所有消息必须是合法 JSON。审批在 http/db 操作前处理。
|
||||
//! 所有消息必须是合法 JSON。
|
||||
//!
|
||||
//! ## 安全边界(授权层)
|
||||
//! **工具级审批在入口统一执行**:高风险工具在 spawn 之前必须获得用户确认,
|
||||
//! 与工具后续输出的消息类型(result/http/db)无关 —— 杜绝 stdio→result 绕过。
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
@@ -24,6 +28,59 @@ struct HttpResponse {
|
||||
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(
|
||||
spec: &ToolSpec,
|
||||
params: &Value,
|
||||
@@ -36,11 +93,32 @@ pub async fn execute_loop(
|
||||
+ Send
|
||||
+ 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 binary_path = spec.path.as_deref().unwrap_or(&spec.tool);
|
||||
|
||||
const MAX_ROUNDS: u32 = 20;
|
||||
|
||||
let mut response: Option<HttpResponse> = None;
|
||||
@@ -51,7 +129,10 @@ pub async fn execute_loop(
|
||||
round += 1;
|
||||
if round > 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 ──
|
||||
let binary = std::path::Path::new(binary_path);
|
||||
let mut cmd = Command::new(&binary);
|
||||
if let Some(c) = &spec.command {
|
||||
cmd.arg("--command").arg(c);
|
||||
}
|
||||
// 按工具 spec 定义的 env 列表,从当前进程环境推送对应变量
|
||||
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.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() {
|
||||
// 执行隔离边界:启用沙箱且工具 profile 非 local 时,通过 docker exec
|
||||
// 在长期容器内执行;否则本地 spawn(env_clear + 白名单注入)。
|
||||
// 沙箱降级(ensure 失败或容器崩溃未恢复)时自动回退本地执行。
|
||||
let use_docker = if crate::tools::sandbox::should_use_docker(&spec.sandbox.profile) {
|
||||
// 执行前确保容器就绪:未运行则重建,重建失败则降级
|
||||
crate::tools::sandbox::ensure_ready_or_degrade("tools").await
|
||||
} else {
|
||||
false
|
||||
};
|
||||
let mut child = match spawn_child(spec, binary_path, use_docker).await {
|
||||
Ok(c) => c,
|
||||
Err(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(Err(e)) => {
|
||||
let _ = child.start_kill();
|
||||
return format!("读取stdout失败: {e}");
|
||||
return ToolRunOutcome::err(ToolErrorKind::StdoutReadFailed, format!("读取stdout失败: {e}"));
|
||||
}
|
||||
Err(_) => {
|
||||
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(Err(e)) => {
|
||||
error!(target: "ias::err","[process] {e}");
|
||||
return format!("进程异常: {e}");
|
||||
return ToolRunOutcome::err(ToolErrorKind::ProcessAbnormal, format!("进程异常: {e}"));
|
||||
}
|
||||
Err(_) => {
|
||||
let _ = child.start_kill();
|
||||
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 {
|
||||
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,
|
||||
Err(_) => {
|
||||
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 elapsed = start.elapsed();
|
||||
info!(target: "ias::tool", "[done] {} 完成, 耗时 {:?}, {} 轮", spec.name, elapsed, round);
|
||||
return content.to_string();
|
||||
return ToolRunOutcome::Ok(content.to_string());
|
||||
}
|
||||
Some("http") => {
|
||||
let method = msg["method"].as_str().unwrap_or("GET");
|
||||
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()
|
||||
.timeout(Duration::from_secs(25))
|
||||
@@ -227,7 +289,9 @@ pub async fn execute_loop(
|
||||
.await
|
||||
}
|
||||
("POST", None) => client.post(url).send().await,
|
||||
(o, _) => return format!("不支持的方法: {o}"),
|
||||
(o, _) => {
|
||||
return ToolRunOutcome::err(ToolErrorKind::UnsupportedMethod, format!("不支持的方法: {o}"))
|
||||
}
|
||||
};
|
||||
match resp {
|
||||
Ok(r) => {
|
||||
@@ -237,32 +301,13 @@ pub async fn execute_loop(
|
||||
}
|
||||
Err(e) => {
|
||||
error!(target: "ias::err","[http] {e}");
|
||||
return format!("HTTP失败: {e}");
|
||||
return ToolRunOutcome::err(ToolErrorKind::HttpFailed, format!("HTTP失败: {e}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
Some("db") => {
|
||||
let op = msg["operation"].as_str().unwrap_or("");
|
||||
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 {
|
||||
"read_memories" => memory.read_for(user_id, None).await,
|
||||
@@ -271,8 +316,7 @@ pub async fn execute_loop(
|
||||
if c.is_empty() {
|
||||
"缺少 content".into()
|
||||
} else {
|
||||
let write_result = memory.write_for(user_id, c).await;
|
||||
write_result
|
||||
memory.write_for(user_id, c).await
|
||||
}
|
||||
}
|
||||
"delete_memory" => {
|
||||
@@ -290,8 +334,79 @@ pub async fn execute_loop(
|
||||
}
|
||||
_ => {
|
||||
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()
|
||||
}
|
||||
|
||||
@@ -6,4 +6,5 @@ pub mod approval;
|
||||
pub mod executor;
|
||||
pub mod parser;
|
||||
pub mod registry;
|
||||
pub mod sandbox;
|
||||
pub mod spec;
|
||||
|
||||
+26
-17
@@ -8,7 +8,7 @@ use std::sync::Arc;
|
||||
|
||||
use crate::context::MemoryStore;
|
||||
use crate::tools::approval::ApprovalManager;
|
||||
use crate::tools::executor;
|
||||
use crate::tools::executor::{self, ToolErrorKind, ToolRunOutcome};
|
||||
use crate::tools::parser;
|
||||
use crate::tools::spec::ToolSpec;
|
||||
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);
|
||||
if path.is_absolute() {
|
||||
return path.to_path_buf();
|
||||
@@ -89,15 +89,22 @@ impl Registry {
|
||||
}
|
||||
|
||||
for mut spec in result.specs {
|
||||
// 解析二进制路径
|
||||
let binary = match &spec.path {
|
||||
Some(p) => {
|
||||
let p = p.trim_start_matches('/');
|
||||
path.join(p)
|
||||
}
|
||||
None => {
|
||||
// 默认: {tool_dir}/target/release/{tool}
|
||||
path.join("target/release").join(&spec.tool)
|
||||
// 优先使用统一产物目录 tools/bin/<tool>,不存在则 fallback 到工具子目录
|
||||
let binary = {
|
||||
let unified = root.join("bin").join(&spec.tool);
|
||||
if unified.exists() {
|
||||
unified
|
||||
} else {
|
||||
match &spec.path {
|
||||
Some(p) => {
|
||||
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
|
||||
+ Sync
|
||||
),
|
||||
) -> String {
|
||||
) -> ToolRunOutcome {
|
||||
let spec = match self.get(name) {
|
||||
Some(s) => s,
|
||||
None => {
|
||||
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) {
|
||||
return err;
|
||||
return ToolRunOutcome::err(ToolErrorKind::MissingParams, err);
|
||||
}
|
||||
executor::execute_loop(
|
||||
spec,
|
||||
@@ -215,20 +222,20 @@ impl RegistryManager {
|
||||
+ Send
|
||||
+ Sync
|
||||
),
|
||||
) -> String {
|
||||
) -> ToolRunOutcome {
|
||||
let spec = {
|
||||
let registry = self.registry.read().await;
|
||||
match registry.get(name) {
|
||||
Some(spec) => spec.clone(),
|
||||
None => {
|
||||
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) {
|
||||
return err;
|
||||
return ToolRunOutcome::err(ToolErrorKind::MissingParams, err);
|
||||
}
|
||||
|
||||
executor::execute_loop(
|
||||
@@ -340,6 +347,7 @@ path: /target/release/mytool
|
||||
required: true,
|
||||
desc: "城市".into(),
|
||||
}],
|
||||
sandbox: Default::default(),
|
||||
};
|
||||
|
||||
let err = validate_required_params(&spec, &serde_json::json!({})).unwrap();
|
||||
@@ -363,6 +371,7 @@ path: /target/release/mytool
|
||||
required: true,
|
||||
desc: "城市".into(),
|
||||
}],
|
||||
sandbox: Default::default(),
|
||||
};
|
||||
|
||||
assert!(
|
||||
|
||||
@@ -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` 包裹:超时 SIGTERM,5 秒后 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>:超时 SIGTERM,5s 后 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()))
|
||||
}
|
||||
@@ -13,6 +13,48 @@ pub struct ParamSpec {
|
||||
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` 的完整映射)
|
||||
///
|
||||
/// ```yaml
|
||||
@@ -56,6 +98,10 @@ pub struct ToolSpec {
|
||||
|
||||
#[serde(default)]
|
||||
pub params: Vec<ParamSpec>,
|
||||
|
||||
/// 沙箱配置(Docker 隔离);默认 profile=local 不启用容器
|
||||
#[serde(default)]
|
||||
pub sandbox: SandboxSpec,
|
||||
}
|
||||
|
||||
fn default_risk_level() -> String {
|
||||
|
||||
@@ -15,3 +15,4 @@ params:
|
||||
- name: city
|
||||
required: false
|
||||
desc: 城市名称,可选。限定地理编码的城市范围
|
||||
|
||||
|
||||
@@ -12,3 +12,4 @@ params:
|
||||
- name: center
|
||||
required: false
|
||||
desc: 地图中心点经纬度坐标,格式"经度,纬度",默认天安门坐标
|
||||
|
||||
|
||||
@@ -18,3 +18,4 @@ params:
|
||||
- name: radius
|
||||
required: false
|
||||
desc: 搜索半径(米),如 500、1000、3000。默认 1000 米
|
||||
|
||||
|
||||
@@ -12,3 +12,4 @@ params:
|
||||
- name: location
|
||||
required: true
|
||||
desc: 经纬度坐标,格式为"经度,纬度",如"116.397,39.908"
|
||||
|
||||
|
||||
@@ -18,3 +18,4 @@ params:
|
||||
- name: type
|
||||
required: false
|
||||
desc: 出行方式:driving(驾车)、walking(步行)、riding(骑行)、transit(公交)。默认 driving
|
||||
|
||||
|
||||
@@ -15,3 +15,4 @@ params:
|
||||
- name: interests
|
||||
required: false
|
||||
desc: 感兴趣的关键词,逗号分隔,如"景点,美食,酒店"。默认"景点,美食"
|
||||
|
||||
|
||||
+5
-4
@@ -1,17 +1,18 @@
|
||||
#!/bin/bash
|
||||
# 构建所有工具
|
||||
# 构建所有工具,产物统一收集到 bin/ 供长期沙箱容器挂载
|
||||
set -e
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
TOOLS=(datetime weather amap web_search fetch_page memories tool_manager)
|
||||
|
||||
mkdir -p bin
|
||||
|
||||
echo "=== 构建工具 ==="
|
||||
for tool in "${TOOLS[@]}"; do
|
||||
echo " $tool..."
|
||||
(cd "$tool" && cargo build --release --quiet)
|
||||
cp "$tool/target/release/$tool" "bin/$tool"
|
||||
done
|
||||
|
||||
echo "=== 完成 ==="
|
||||
for tool in "${TOOLS[@]}"; do
|
||||
ls -lh "$tool/target/release/$tool" 2>/dev/null
|
||||
done
|
||||
ls -lh bin/*
|
||||
|
||||
@@ -6,3 +6,4 @@ path: /target/release/datetime
|
||||
risk_level: low
|
||||
timeout_secs: 3
|
||||
params: []
|
||||
|
||||
|
||||
@@ -9,3 +9,4 @@ params:
|
||||
- name: url
|
||||
required: true
|
||||
desc: 网页URL
|
||||
|
||||
|
||||
@@ -12,3 +12,4 @@ params:
|
||||
- name: memory_id
|
||||
required: true
|
||||
desc: 记忆ID
|
||||
|
||||
|
||||
@@ -12,3 +12,4 @@ params:
|
||||
- name: memory_id
|
||||
required: true
|
||||
desc: 记忆的数字ID,从 read_memories 返回结果中获取
|
||||
|
||||
|
||||
@@ -9,3 +9,4 @@ timeout_secs: 3
|
||||
env:
|
||||
- IAS_MEMORY_DIR
|
||||
params:
|
||||
|
||||
|
||||
@@ -15,3 +15,4 @@ params:
|
||||
- name: content
|
||||
required: true
|
||||
desc: 更新后的新内容,覆盖原记忆
|
||||
|
||||
|
||||
@@ -12,3 +12,4 @@ params:
|
||||
- name: content
|
||||
required: true
|
||||
desc: 要记录的记忆内容,用自然语言描述。如"用户住在合肥包河区联投新安里A区"
|
||||
|
||||
|
||||
Generated
+105
@@ -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"
|
||||
@@ -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 目录
|
||||
@@ -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
|
||||
path: /target/release/tool_manager
|
||||
risk_level: high
|
||||
timeout_secs: 120
|
||||
timeout_secs: 300
|
||||
env:
|
||||
- HOME
|
||||
- CARGO_HOME
|
||||
- RUSTUP_HOME
|
||||
params:
|
||||
- name: action
|
||||
required: true
|
||||
@@ -23,7 +27,7 @@ params:
|
||||
desc: 参数列表数组,每项包含 name、required、desc
|
||||
- name: risk_level
|
||||
required: false
|
||||
desc: 新工具风险级别,low 或 high,默认 low
|
||||
desc: 已弃用,tool_manager 创建/更新的工具一律强制 high risk,传入值将被忽略
|
||||
- name: timeout_secs
|
||||
required: false
|
||||
desc: 新工具超时时间秒数,默认 30
|
||||
|
||||
+169
-30
@@ -1,10 +1,22 @@
|
||||
//! tool_manager — 受控创建、修改、构建本地工具
|
||||
//!
|
||||
//! 安全约束:
|
||||
//! - 不允许操作保留工具名(含 `tool_manager` 自身),防止自改后门
|
||||
//! - 由本工具创建/更新的工具一律 `risk_level: high`,忽略调用方传入值
|
||||
//! - 所有调用追加写入 `tools/.tool_audit.log` 审计日志
|
||||
//! - 构建使用 `timeout` 限制时长,失败时清理残留二进制,避免调用旧版本"假成功"
|
||||
|
||||
use serde_json::{Value, json};
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
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() {
|
||||
let input: Value = match serde_json::from_reader(std::io::stdin()) {
|
||||
Ok(v) => v,
|
||||
@@ -12,13 +24,20 @@ fn main() {
|
||||
};
|
||||
|
||||
let params = &input["params"];
|
||||
let user_id = input["user_id"].as_str().unwrap_or("unknown");
|
||||
let action = params["action"].as_str().unwrap_or("");
|
||||
let tool_name = params["tool_name"].as_str().unwrap_or("");
|
||||
|
||||
let output = match run(action, tool_name, params) {
|
||||
Ok(msg) => msg,
|
||||
Err(e) => format!("工具管理失败: {e}"),
|
||||
let (output, success) = match run(action, tool_name, params) {
|
||||
Ok(msg) => (msg, true),
|
||||
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);
|
||||
}
|
||||
|
||||
@@ -38,7 +57,7 @@ fn run(action: &str, tool_name: &str, params: &Value) -> Result<String, String>
|
||||
write_tool_project(&tool_dir, tool_name, params)?;
|
||||
build_tool(&tool_dir, tool_name)?;
|
||||
Ok(format!(
|
||||
"已创建并构建工具 {tool_name}。可通过 query_capabilities 查看并用 call_capability 调用。"
|
||||
"已创建并构建工具 {tool_name}(risk_level 强制为 high)。可通过 query_capabilities 查看并用 call_capability 调用。"
|
||||
))
|
||||
}
|
||||
"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)?;
|
||||
build_tool(&tool_dir, tool_name)?;
|
||||
Ok(format!(
|
||||
"已修改并构建工具 {tool_name}。注册表重载后即可使用最新版本。"
|
||||
"已修改并构建工具 {tool_name}(risk_level 强制为 high)。注册表重载后即可使用最新版本。"
|
||||
))
|
||||
}
|
||||
"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())?;
|
||||
}
|
||||
|
||||
// spec 一旦由 tool_manager 管理,risk_level 恒为 high
|
||||
let should_update_spec = params.get("description").is_some()
|
||||
|| params.get("params").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> {
|
||||
let output = Command::new("cargo")
|
||||
.arg("build")
|
||||
.arg("--release")
|
||||
.current_dir(tool_dir)
|
||||
.output()
|
||||
.map_err(|e| format!("启动 cargo 失败: {e}"))?;
|
||||
let binary = tool_dir.join("target/release").join(tool_name);
|
||||
let use_builder = std::env::var("IAS_DOCKER_BUILDER")
|
||||
.ok()
|
||||
.as_deref()
|
||||
== Some("1");
|
||||
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() {
|
||||
// 构建失败:删除可能残留的旧二进制,避免 reload 后仍调用旧版本造成"假成功"
|
||||
let _ = fs::remove_file(&binary);
|
||||
let code = output.status.code().unwrap_or(-1);
|
||||
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() {
|
||||
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(())
|
||||
}
|
||||
|
||||
/// 定位 tools 根目录。
|
||||
///
|
||||
/// 产物可能位于两处之一:
|
||||
/// - `tools/bin/tool_manager`(统一产物目录,registry 优先使用)
|
||||
/// - `tools/tool_manager/target/release/tool_manager`(构建产物)
|
||||
///
|
||||
/// 因此不依赖固定向上层数,而是从 exe 向上找第一个“含 bin/ 子目录且至少含一个
|
||||
/// 工具子目录”的祖先,即为 tools 根。对两种位置均正确。
|
||||
fn locate_tools_root() -> Result<PathBuf, String> {
|
||||
let exe = std::env::current_exe().map_err(|e| e.to_string())?;
|
||||
let tool_dir = exe
|
||||
.parent()
|
||||
.and_then(Path::parent)
|
||||
.and_then(Path::parent)
|
||||
.ok_or_else(|| "无法定位 tool_manager 目录".to_string())?;
|
||||
tool_dir
|
||||
.parent()
|
||||
.map(Path::to_path_buf)
|
||||
.ok_or_else(|| "无法定位 tools 根目录".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 cargo_toml(tool_name: &str) -> String {
|
||||
@@ -171,14 +264,13 @@ fn main() {{
|
||||
}
|
||||
|
||||
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"]
|
||||
.as_str()
|
||||
.unwrap_or("由 tool_manager 创建的本地工具");
|
||||
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);
|
||||
.unwrap_or("由 tool_manager 创建的本地工具(自动标记 high risk)");
|
||||
|
||||
let mut out = format!(
|
||||
"name: {tool_name}\n\
|
||||
@@ -211,6 +303,10 @@ fn spec_yaml(tool_name: &str, params: &Value) -> String {
|
||||
out.push_str("params: []\n");
|
||||
}
|
||||
|
||||
// 安全约束:tool_manager 产出的工具不受信任(含 LLM 提供的 rust_code),
|
||||
// 默认进 stdio 沙箱容器(无网络、只读根),防止数据外泄/越权文件读
|
||||
out.push_str("sandbox:\n profile: stdio\n");
|
||||
|
||||
out
|
||||
}
|
||||
|
||||
@@ -218,8 +314,10 @@ fn validate_tool_name(name: &str) -> Result<(), String> {
|
||||
if name.is_empty() {
|
||||
return Err("缺少 tool_name".to_string());
|
||||
}
|
||||
if matches!(name, "target" | "specs" | "src" | "." | "..") {
|
||||
return Err("tool_name 使用了保留名称".to_string());
|
||||
if RESERVED_NAMES.contains(&name) {
|
||||
return Err(format!(
|
||||
"tool_name '{name}' 是保留名称,禁止 create/update/build"
|
||||
));
|
||||
}
|
||||
if !name
|
||||
.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) {
|
||||
println!("{}", json!({"type":"result","content":content}));
|
||||
}
|
||||
|
||||
@@ -12,3 +12,4 @@ params:
|
||||
- name: location
|
||||
required: true
|
||||
desc: 城市名称,如"北京"、"上海"、"合肥"。支持地级市和县级市
|
||||
|
||||
|
||||
@@ -12,3 +12,4 @@ params:
|
||||
- name: location
|
||||
required: true
|
||||
desc: 城市名称,如"北京"、"上海"、"合肥"。支持地级市和县级市
|
||||
|
||||
|
||||
@@ -12,3 +12,4 @@ params:
|
||||
- name: location
|
||||
required: true
|
||||
desc: 城市名称,如"北京"、"上海"、"合肥"。支持地级市和县级市
|
||||
|
||||
|
||||
@@ -13,3 +13,4 @@ params:
|
||||
desc: 搜索关键词,如"Rust 2024 edition 新特性"、"今天比特币价格"
|
||||
- name: max_results
|
||||
desc: 返回的最大搜索结果数量(1-20),默认5条
|
||||
|
||||
|
||||
Reference in New Issue
Block a user