工具系统重构与定时任务管理

1. 工具抽象体系
   - 新增 Tool trait + ToolRegistry 注册表
   - 三种工具类型:ApiTool / StdioTool / ShellTool
   - 全局 REGISTRY (LazyLock) 作为主入口
   - BuiltinRegistry 简化为兼容层

2. 定时任务管理工具
   - 新增 manage_scheduled_tasks 工具(ApiTool)
   - 支持 list/add/update/delete/toggle 五种操作
   - 通过 LazyLock 延迟初始化数据库连接池

3. 定时任务分流
   - 新增 task_type 字段(model / system)
   - 模型可自由管理 model 任务
   - system 任务只读,修改操作被拒绝

4. CLI 支持
   - 新增 ias task 子命令(list/add/update/delete/toggle)

5. 清理
   - 删除废弃的 ias-auth 独立二进制
   - SkillResult 增加 Deserialize derive
   - 所有调用点更新为 crate::tools::xxx 新入口
This commit is contained in:
2026-06-10 16:00:47 +08:00
parent f7d0781090
commit 5a54368437
17 changed files with 1271 additions and 257 deletions
-3
View File
@@ -52,6 +52,3 @@ scraper = "0.27.0" # HTML 解析(f
name = "ias"
path = "src/main.rs"
[[bin]]
name = "ias-auth"
path = "src/bin/ias-auth.rs"
@@ -0,0 +1,6 @@
-- 定时任务增加 task_type 字段:system(系统任务)或 model(模型任务)
-- 模型只能管理 model 类型的任务,不能修改 system 任务
ALTER TABLE scheduled_tasks ADD COLUMN IF NOT EXISTS task_type TEXT NOT NULL DEFAULT 'model';
CREATE INDEX IF NOT EXISTS idx_scheduled_tasks_type
ON scheduled_tasks (task_type, enabled, next_run_at);
-112
View File
@@ -1,112 +0,0 @@
/// ## ias-auth —— 独立 Token 生成器(辅助二进制)
///
/// 这是一个独立的二进制文件,不依赖主程序的任何模块。
/// 专供 JSON 工具定义通过 `${cmd:ias-auth gen-qweather-jwt}` 调用,
/// 用于生成和风天气 API 的 Ed25519 JWT Token。
///
/// ### 设计原因
/// 独立编译有两个好处:
/// 1. 编译速度快(没有主程序的庞大依赖链)
/// 2. 在 JSON 工具定义中通过 shell 命令调用的开销很小
///
/// ### 用法
/// ```bash
/// ias-auth gen-qweather-jwt
/// ```
use std::process::exit;
fn main() {
let args: Vec<String> = std::env::args().collect();
if args.len() < 2 {
eprintln!("用法: ias-auth <command>");
eprintln!("命令:");
eprintln!(" gen-qweather-jwt 生成和风天气 Ed25519 JWT Token");
exit(1);
}
match args[1].as_str() {
"gen-qweather-jwt" => gen_qweather_jwt(),
_ => {
eprintln!("未知命令: {}", args[1]);
eprintln!("支持的命令: gen-qweather-jwt");
exit(1);
}
}
}
fn gen_qweather_jwt() {
match generate_jwt() {
Ok(token) => println!("{}", token),
Err(e) => {
eprintln!("{}", e);
exit(1);
}
}
}
// ─── 和风天气 Ed25519 JWT 生成 ───
// 从 iAs 的 tools::builtins::weather 提取,保持独立
use base64::Engine;
use chrono::Utc;
use ed25519_dalek::pkcs8::DecodePrivateKey;
use ed25519_dalek::{Signer, SigningKey};
fn base64url(data: &[u8]) -> String {
base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(data)
}
fn generate_jwt() -> Result<String, String> {
let key_file = std::env::var("QWEATHER_JWT_PRIVATE_KEY_FILE")
.unwrap_or_else(|_| "qweather/ed25519-private.pem".into());
let pem = std::fs::read_to_string(&key_file)
.map_err(|e| format!("读取密钥文件 {} 失败: {}", key_file, e))?;
// 手动解析 PEM: 去掉头尾,base64 解码得到 DER
let der = pem
.lines()
.filter(|l| !l.starts_with("-----"))
.collect::<Vec<_>>()
.join("");
let der_bytes = base64::engine::general_purpose::STANDARD
.decode(&der)
.map_err(|e| format!("Base64 解码密钥失败: {e}"))?;
let signing_key = SigningKey::from_pkcs8_der(&der_bytes)
.map_err(|e| format!("解析 Ed25519 密钥失败: {e}"))?;
let key_id = std::env::var("QWEATHER_JWT_KEY_ID").unwrap_or_default();
let project_id = std::env::var("QWEATHER_JWT_PROJECT_ID").unwrap_or_default();
let ttl: i64 = std::env::var("QWEATHER_JWT_TTL_SECONDS")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(3600);
let now = Utc::now().timestamp();
// Header
let header = serde_json::json!({"alg": "EdDSA", "kid": key_id});
let header_b64 = base64url(header.to_string().as_bytes());
// Payload
let payload = serde_json::json!({
"sub": project_id,
"iat": now - 30,
"exp": now + ttl,
});
let payload_b64 = base64url(payload.to_string().as_bytes());
// Sign
let message = format!("{}.{}", header_b64, payload_b64);
let signature = signing_key
.try_sign(message.as_bytes())
.map_err(|e| format!("JWT 签名失败: {e}"))?;
Ok(format!(
"{}.{}.{}",
header_b64,
payload_b64,
base64url(&signature.to_bytes())
))
}
+106
View File
@@ -142,6 +142,19 @@ pub enum Commands {
sock: String,
},
/// 管理定时任务(增删改查)
///
/// 需要 PostgreSQL 数据库。
///
/// 示例:
/// ias task list
/// ias task add --name 每日天气 --user wx_xxx --command "curl -s wttr.in/Beijing?format=3" --interval 86400
/// ias task update <uuid> --name "新名称" --interval 3600
/// ias task delete <uuid>
/// ias task toggle <uuid>
#[command(subcommand)]
ScheduledTask(ScheduledTaskAction),
/// 调用内置工具(无需登录,独立运行)
///
/// 可直接在终端调用天气、搜索、备忘录、日期时间、高德地图等工具。
@@ -152,6 +165,99 @@ pub enum Commands {
Tool(ToolCommand),
}
// ─── 定时任务子命令 ───
#[derive(Subcommand, Debug, Clone)]
pub enum ScheduledTaskAction {
/// 列出所有定时任务
List,
/// 添加定时任务
///
/// 示例:
/// ias task add --name 每日天气 --user wx_xxx --command "curl -s wttr.in/Beijing?format=3" --interval 86400
Add {
/// 任务名称(唯一)
#[arg(long)]
name: String,
/// 所属用户 ID
#[arg(long)]
user: String,
/// 要执行的 shell 命令
#[arg(long)]
command: String,
/// 执行间隔(秒)
#[arg(long)]
interval: u32,
/// 任务描述
#[arg(long)]
description: Option<String>,
/// Shell 解释器(默认 /bin/bash
#[arg(long)]
shell: Option<String>,
/// 工作目录
#[arg(long)]
cwd: Option<String>,
},
/// 更新定时任务
///
/// 示例:
/// ias task update <uuid> --name "新名称" --interval 3600
Update {
/// 任务 IDUUID
id: String,
/// 新名称
#[arg(long)]
name: Option<String>,
/// 新命令
#[arg(long)]
command: Option<String>,
/// 新执行间隔(秒)
#[arg(long)]
interval: Option<u32>,
/// 新描述
#[arg(long)]
description: Option<String>,
/// 新 Shell 解释器
#[arg(long)]
shell: Option<String>,
/// 新工作目录
#[arg(long)]
cwd: Option<String>,
},
/// 删除定时任务
///
/// 示例:
/// ias task delete <uuid>
Delete {
/// 任务 IDUUID
id: String,
},
/// 启用/禁用定时任务
///
/// 示例:
/// ias task toggle <uuid>
Toggle {
/// 任务 IDUUID
id: String,
},
}
// ─── 工具子命令 ───
#[derive(Subcommand, Debug, Clone)]
+6 -6
View File
@@ -549,7 +549,7 @@ async fn tool_consumer_loop(
info!("Tool Consumer: user={} tool={}", user_id, tool_name);
// 高风险工具 → 审批(除非已被用户批准)
if !approved && crate::tools::builtin::BuiltinRegistry::is_high_risk(&tool_name) {
if !approved && crate::tools::is_high_risk(&tool_name) {
info!("高风险工具 {} 需要审批", tool_name);
ctx.approval_ctx.lock().await.insert(
user_id.clone(),
@@ -620,7 +620,7 @@ async fn execute_tool(name: &str, arguments: &str, ctx: &DaemonCtx, user_id: &st
.collect::<Vec<_>>().join("\n---\n");
}
"query_capabilities" => {
let specs = crate::tools::builtin::BuiltinRegistry::specs();
let specs = crate::tools::specs();
return crate::tools::build_capability_guide(&specs);
}
"call_capability" => {
@@ -656,13 +656,13 @@ async fn execute_tool(name: &str, arguments: &str, ctx: &DaemonCtx, user_id: &st
.collect::<Vec<_>>().join("\n---\n");
}
"query_capabilities" => {
let specs = crate::tools::builtin::BuiltinRegistry::specs();
let specs = crate::tools::specs();
return crate::tools::build_capability_guide(&specs);
}
_ => {}
}
// 内置工具注册表
match crate::tools::builtin::BuiltinRegistry::execute(&target_name, &target_args).await {
match crate::tools::execute(&target_name, &target_args).await {
Some(r) => return r.output,
None => return format!("未知工具: {}", target_name),
}
@@ -671,7 +671,7 @@ async fn execute_tool(name: &str, arguments: &str, ctx: &DaemonCtx, user_id: &st
}
// 内置工具注册表
match crate::tools::builtin::BuiltinRegistry::execute(name, arguments).await {
match crate::tools::execute(name, arguments).await {
Some(r) => r.output,
None => format!("未知工具: {}", name),
}
@@ -795,7 +795,7 @@ fn build_tools_list() -> Vec<serde_json::Value> {
}
}));
let specs = crate::tools::builtin::BuiltinRegistry::specs();
let specs = crate::tools::specs();
for spec in &specs {
list.push(serde_json::json!({
"type": "function",
+160
View File
@@ -243,3 +243,163 @@ pub async fn load_summaries(pool: &PgPool, user_id: &str, limit: i64) -> Result<
.map_err(|e| format!("查询摘要失败: {}", e))?;
Ok(rows.into_iter().map(|(t,)| t).collect())
}
// ─── 定时任务 ───
#[derive(Debug, Clone, sqlx::FromRow, Serialize, Deserialize)]
pub struct ScheduledTask {
pub id: uuid::Uuid,
pub name: String,
pub description: String,
pub user_id: Option<String>,
pub command: String,
pub shell: String,
pub cwd: String,
pub interval_seconds: i32,
pub enabled: bool,
pub next_run_at: DateTime<Utc>,
pub last_run_at: Option<DateTime<Utc>>,
pub last_status: Option<String>,
pub task_type: String,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
}
/// 列出所有定时任务
pub async fn list_scheduled_tasks(pool: &PgPool) -> Result<Vec<ScheduledTask>, String> {
let tasks = sqlx::query_as::<_, ScheduledTask>(
r#"
SELECT id, name, description, user_id, command, shell, cwd,
interval_seconds, enabled, next_run_at, last_run_at,
last_status, task_type, created_at, updated_at
FROM scheduled_tasks
ORDER BY next_run_at ASC
"#,
)
.fetch_all(pool)
.await
.map_err(|e| format!("查询定时任务失败: {}", e))?;
Ok(tasks)
}
/// 添加定时任务
pub async fn add_scheduled_task(
pool: &PgPool,
name: &str,
user_id: &str,
command: &str,
interval_seconds: i32,
description: &str,
shell: &str,
cwd: &str,
task_type: &str,
) -> Result<uuid::Uuid, String> {
let id = uuid::Uuid::new_v4();
sqlx::query(
r#"
INSERT INTO scheduled_tasks (id, name, description, user_id, command, shell, cwd, interval_seconds, task_type)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
"#,
)
.bind(id)
.bind(name)
.bind(description)
.bind(user_id)
.bind(command)
.bind(shell)
.bind(cwd)
.bind(interval_seconds)
.bind(task_type)
.execute(pool)
.await
.map_err(|e| format!("添加定时任务失败: {}", e))?;
Ok(id)
}
/// 更新定时任务(只更新非 None 的字段)
pub async fn update_scheduled_task(
pool: &PgPool,
id: &uuid::Uuid,
name: Option<&str>,
command: Option<&str>,
interval_seconds: Option<i32>,
description: Option<&str>,
shell: Option<&str>,
cwd: Option<&str>,
) -> Result<bool, String> {
let mut sets = Vec::new();
let mut idx = 1u32;
if name.is_some() {
sets.push(format!("name = ${}", idx)); idx += 1;
}
if command.is_some() {
sets.push(format!("command = ${}", idx)); idx += 1;
}
if interval_seconds.is_some() {
sets.push(format!("interval_seconds = ${}", idx)); idx += 1;
}
if description.is_some() {
sets.push(format!("description = ${}", idx)); idx += 1;
}
if shell.is_some() {
sets.push(format!("shell = ${}", idx)); idx += 1;
}
if cwd.is_some() {
sets.push(format!("cwd = ${}", idx)); idx += 1;
}
if sets.is_empty() {
return Ok(false);
}
sets.push("updated_at = NOW()".to_string());
let sql = format!(
"UPDATE scheduled_tasks SET {} WHERE id = ${}",
sets.join(", "),
idx
);
let mut query = sqlx::query(&sql).bind(id);
if let Some(v) = name { query = query.bind(v); }
if let Some(v) = command { query = query.bind(v); }
if let Some(v) = interval_seconds { query = query.bind(v); }
if let Some(v) = description { query = query.bind(v); }
if let Some(v) = shell { query = query.bind(v); }
if let Some(v) = cwd { query = query.bind(v); }
let rows = query
.execute(pool)
.await
.map_err(|e| format!("更新定时任务失败: {}", e))?;
Ok(rows.rows_affected() > 0)
}
/// 删除定时任务
pub async fn delete_scheduled_task(pool: &PgPool, id: &uuid::Uuid) -> Result<bool, String> {
let rows = sqlx::query("DELETE FROM scheduled_tasks WHERE id = $1")
.bind(id)
.execute(pool)
.await
.map_err(|e| format!("删除定时任务失败: {}", e))?;
Ok(rows.rows_affected() > 0)
}
/// 切换定时任务启用状态
pub async fn toggle_scheduled_task(pool: &PgPool, id: &uuid::Uuid) -> Result<(bool, bool), String> {
let row = sqlx::query_as::<_, (bool,)>(
"UPDATE scheduled_tasks SET enabled = NOT enabled, updated_at = NOW() WHERE id = $1 RETURNING enabled",
)
.bind(id)
.fetch_optional(pool)
.await
.map_err(|e| format!("切换定时任务状态失败: {}", e))?;
match row {
Some((enabled,)) => Ok((true, enabled)),
None => Ok((false, false)),
}
}
+160 -6
View File
@@ -8,7 +8,7 @@
//! ├── ias listen --llm → Daemon 模式(主运行模式)
//! │ └── daemon.rs: 长轮询 → MessageQueue → 3 个 tokio 消费者
//! │ ├── LLM Consumer → Conversation.chat_with_tools() → DeepSeek API
//! │ ├── Tool Consumer → BuiltinRegistry.execute() / ApprovalManager
//! │ ├── Tool Consumer → tools::execute() / ApprovalManager
//! │ └── Send Consumer → WeChatClient.send_message()
//! ├── ias service → 等同 listen --llm 但启用文件日志
//! ├── ias daemon → 守护进程入口(等同 listen --llm
@@ -45,7 +45,7 @@ mod wechat;
mod worker;
use clap::Parser;
use cli::{AmapAction, Cli, Commands, MemosAction, ToolCommand};
use cli::{AmapAction, Cli, Commands, MemosAction, ScheduledTaskAction, ToolCommand};
use context::MemoryStore;
use db::Database;
use llm::{ChatResult, Conversation, ConversationConfig, DEFAULT_SYSTEM_PROMPT, ToolExecutor, Usage};
@@ -144,6 +144,9 @@ async fn main() {
error!("Worker 失败: {}", e);
}
}
Commands::ScheduledTask(action) => {
cmd_scheduled_task(action, &database).await;
}
Commands::Tool(..) => unreachable!("Tool 命令已提前处理"),
}
}
@@ -404,7 +407,7 @@ async fn cmd_listen(
// query_capabilities: 列出所有内置工具
if n == "query_capabilities" {
let specs = tools::builtin::BuiltinRegistry::specs();
let specs = tools::specs();
return Ok(tools::build_capability_guide(&specs));
}
@@ -424,8 +427,8 @@ async fn cmd_listen(
};
// 内置工具
if tools::builtin::BuiltinRegistry::is_builtin(&target_name) {
if tools::builtin::BuiltinRegistry::is_high_risk(&target_name) {
if tools::is_builtin(&target_name) {
if tools::is_high_risk(&target_name) {
let mut ctx = ExecutionContext::new(&user_id);
ctx = ctx.with_approval(approval.clone());
ctx = ctx.with_wechat(send.clone());
@@ -437,7 +440,7 @@ async fn cmd_listen(
}
}
if let Some(result) =
tools::builtin::BuiltinRegistry::execute(&target_name, &target_args).await
tools::execute(&target_name, &target_args).await
{
return Ok(result.output);
}
@@ -900,6 +903,157 @@ async fn builtin_approve(ctx: &ExecutionContext, name: &str) -> Result<bool, Str
}
}
// ─── 定时任务 CLI 命令 ───
async fn cmd_scheduled_task(action: ScheduledTaskAction, database: &Option<Arc<Database>>) {
let db = match database {
Some(d) => d,
None => {
println!("❌ 数据库未配置,无法管理定时任务");
return;
}
};
match action {
ScheduledTaskAction::List => {
match db::models::list_scheduled_tasks(db.pool()).await {
Ok(tasks) => {
if tasks.is_empty() {
println!("📋 没有定时任务");
return;
}
println!("📋 定时任务列表:");
println!("{:=<80}", "");
for task in &tasks {
let status = if task.enabled { "🟢" } else { "🔴" };
let last = task.last_run_at
.map(|t| t.format("%Y-%m-%d %H:%M:%S").to_string())
.unwrap_or_else(|| "-".to_string());
let next = task.next_run_at.format("%Y-%m-%d %H:%M:%S");
let interval = if task.interval_seconds >= 86400 {
format!("{}", task.interval_seconds / 86400)
} else if task.interval_seconds >= 3600 {
format!("{}小时", task.interval_seconds / 3600)
} else {
format!("{}", task.interval_seconds)
};
println!("{} {}", status, task.name);
println!(" ID: {}", task.id);
println!(" 用户: {}", task.user_id.as_deref().unwrap_or("-"));
println!(" 命令: {}", task.command);
println!(" 间隔: {}", interval);
println!(" 上次执行: {}", last);
println!(" 下次执行: {}", next);
if let Some(ref s) = task.last_status {
println!(" 上次状态: {:.100}", s);
}
println!("{:=<80}", "");
}
}
Err(e) => println!("❌ 查询失败: {}", e),
}
}
ScheduledTaskAction::Add {
name,
user,
command,
interval,
description,
shell,
cwd,
} => {
let desc = description.unwrap_or_default();
let sh = shell.unwrap_or_else(|| "/bin/bash".to_string());
let wd = cwd.unwrap_or_default();
match db::models::add_scheduled_task(
db.pool(),
&name,
&user,
&command,
interval as i32,
&desc,
&sh,
&wd,
"model",
)
.await
{
Ok(id) => {
println!("✅ 定时任务已添加");
println!(" ID: {}", id);
println!(" 名称: {}", name);
println!(" 间隔: {}", interval);
}
Err(e) => println!("❌ 添加失败: {}", e),
}
}
ScheduledTaskAction::Update {
id,
name,
command,
interval,
description,
shell,
cwd,
} => {
let uuid = match uuid::Uuid::parse_str(&id) {
Ok(u) => u,
Err(e) => {
println!("❌ 无效的 UUID: {}", e);
return;
}
};
match db::models::update_scheduled_task(
db.pool(),
&uuid,
name.as_deref(),
command.as_deref(),
interval.map(|v| v as i32),
description.as_deref(),
shell.as_deref(),
cwd.as_deref(),
)
.await
{
Ok(true) => println!("✅ 定时任务已更新"),
Ok(false) => println!("⚠️ 未找到该任务或无字段更新"),
Err(e) => println!("❌ 更新失败: {}", e),
}
}
ScheduledTaskAction::Delete { id } => {
let uuid = match uuid::Uuid::parse_str(&id) {
Ok(u) => u,
Err(e) => {
println!("❌ 无效的 UUID: {}", e);
return;
}
};
match db::models::delete_scheduled_task(db.pool(), &uuid).await {
Ok(true) => println!("✅ 定时任务已删除"),
Ok(false) => println!("⚠️ 未找到该任务"),
Err(e) => println!("❌ 删除失败: {}", e),
}
}
ScheduledTaskAction::Toggle { id } => {
let uuid = match uuid::Uuid::parse_str(&id) {
Ok(u) => u,
Err(e) => {
println!("❌ 无效的 UUID: {}", e);
return;
}
};
match db::models::toggle_scheduled_task(db.pool(), &uuid).await {
Ok((true, enabled)) => {
let status = if enabled { "🟢 已启用" } else { "🔴 已禁用" };
println!("✅ 定时任务状态已切换: {}", status);
}
Ok((false, _)) => println!("⚠️ 未找到该任务"),
Err(e) => println!("❌ 切换失败: {}", e),
}
}
}
}
// ─── 工具 CLI 命令 ───
async fn cmd_tool(cmd: ToolCommand) {
+45
View File
@@ -0,0 +1,45 @@
//! ## API 工具 —— 通过 Rust 代码实现的工具
//!
//! 所有通过 Rust 代码直接实现的工具(HTTP API 调用、本地计算、文件 I/O 等)
//! 都包装为 `ApiTool`。它是目前最常用的工具类型。
use async_trait::async_trait;
use super::tool::{Tool, ToolHandler};
use super::types::{SkillResult, SkillSpec};
/// ## API 工具
///
/// 包装一个异步闭包作为工具实现。
/// 适用于所有通过 Rust 代码直接实现的工具。
///
/// ### 示例
/// ```rust,ignore
/// let tool = ApiTool::new(
/// datetime::spec(),
/// Arc::new(|params| Box::pin(async move {
/// Ok(SkillResult::ok("done"))
/// })),
/// );
/// ```
pub struct ApiTool {
spec: SkillSpec,
handler: ToolHandler,
}
impl ApiTool {
pub fn new(spec: SkillSpec, handler: ToolHandler) -> Self {
Self { spec, handler }
}
}
#[async_trait]
impl Tool for ApiTool {
fn spec(&self) -> &SkillSpec {
&self.spec
}
async fn execute(&self, params: serde_json::Value) -> SkillResult {
(self.handler)(params).await
}
}
+17 -64
View File
@@ -1,83 +1,36 @@
//! ## 内置工具注册表 —— 工具路由中心
//! ## 内置工具注册表 —— 旧接口兼容层
//!
//! 连接 LLM 的元工具调用和实际 Rust 实现的工具函数
//! 所有内置工具在此统一注册,通过名称路由到对应实现
//!
//! ### 职责
//! 1. `execute(name, args)` — 按名称路由到对应的 builtin 工具执行
//! 2. `specs()` — 返回所有内置工具的 SkillSpec(给元工具 `query_capabilities` 使用)
//! 3. `is_high_risk(name)` — 判断工具是否需要审批
//! 4. `is_builtin(name)` — 判断名称是否对应一个已注册的内置工具
//! 保留 `BuiltinRegistry` 结构体以保持向后兼容
//! 所有方法委托给 `crate::tools::` 模块级函数(新主入口)
use crate::tools::types::{RiskLevel, SkillResult};
use crate::tools::types::{SkillResult, SkillSpec};
/// ## 内置工具注册表
/// ## 内置工具注册表(旧接口兼容)
///
/// 连接 LLM 的元工具调用和实际 Rust 实现的工具函数
///
/// ### 职责
/// 1. `execute(name, args)` — 按名称路由到对应的 builtin 工具执行
/// 2. `specs()` — 返回所有内置工具的 SkillSpec(给元工具 `query_capabilities` 使用)
/// 3. `is_high_risk(name)` — 判断工具是否需要审批
///
/// ### 当前内置工具
/// * `get_current_datetime` — 获取当前北京时间
/// * `manage_memos` — 管理备忘录(CRUD
/// * `query_weather` — 查询天气(和风天气 API)
/// * `web_search` — 联网搜索(Tavily API
/// * `fetch_page` — 抓取网页内容
/// * `amap_*` — 高德地图系列(POI搜索、地理编码、路径规划等)
/// 保持与现有代码兼容的静态方法接口
/// 所有方法委托给 `crate::tools::` 模块级函数。
#[allow(dead_code)]
pub struct BuiltinRegistry;
#[allow(dead_code)]
impl BuiltinRegistry {
/// 按名称执行工具
pub async fn execute(name: &str, args_json: &str) -> Option<SkillResult> {
match name {
"get_current_datetime" => Some(super::builtins::datetime::execute()),
"manage_memos" => Some(super::builtins::memos::execute(args_json).await),
"query_weather" => {
let params: serde_json::Value =
serde_json::from_str(args_json).unwrap_or_default();
Some(super::builtins::weather::execute(params).await)
}
"web_search" => {
let params: serde_json::Value =
serde_json::from_str(args_json).unwrap_or_default();
Some(super::builtins::web_search::execute(params).await)
}
"fetch_page" => {
let params: serde_json::Value =
serde_json::from_str(args_json).unwrap_or_default();
Some(super::builtins::fetch_page::execute(params).await)
}
n if n.starts_with("amap_") => {
let params: serde_json::Value =
serde_json::from_str(args_json).unwrap_or_default();
super::builtins::amap::execute(n, params).await
}
_ => None,
}
crate::tools::execute(name, args_json).await
}
/// 返回所有内置工具的 spec 列表
pub fn specs() -> Vec<super::types::SkillSpec> {
let mut specs = vec![
super::builtins::datetime::spec(),
super::builtins::memos::spec(),
super::builtins::weather::spec(),
super::builtins::web_search::spec(),
super::builtins::fetch_page::spec(),
];
specs.extend(super::builtins::amap::specs());
specs
pub fn specs() -> Vec<SkillSpec> {
crate::tools::specs()
}
/// 判断工具是否为高风险
pub fn is_high_risk(name: &str) -> bool {
Self::specs()
.iter()
.any(|s| s.name == name && s.risk_level == RiskLevel::High)
crate::tools::is_high_risk(name)
}
/// 判断名称是否对应一个已注册的内置工具
pub fn is_builtin(name: &str) -> bool {
Self::specs().iter().any(|s| s.name == name)
crate::tools::is_builtin(name)
}
}
+1
View File
@@ -18,5 +18,6 @@ pub mod amap;
pub mod datetime;
pub mod fetch_page;
pub mod memos;
pub mod scheduled_task;
pub mod weather;
pub mod web_search;
+278
View File
@@ -0,0 +1,278 @@
//! ## 定时任务管理工具 —— 增删改查定时任务
//!
//! 通过 PostgreSQL 数据库管理 `scheduled_tasks` 表。
//! 支持 list / add / update / delete / toggle 五种操作。
//!
//! ### 任务类型
//! - `model` — 模型任务,LLM 可以自由增删改查
//! - `system` — 系统任务,LLM 只能查看,不能修改/删除
use std::sync::LazyLock;
use crate::tools::types::{RiskLevel, SkillResult, SkillSpec};
/// 延迟初始化的数据库连接池
static POOL: LazyLock<Option<sqlx::PgPool>> = LazyLock::new(|| {
let url = std::env::var("DATABASE_URL").ok()?;
let pool = sqlx::PgPool::connect_lazy(&url).ok()?;
Some(pool)
});
/// 返回该工具的元数据描述
pub fn spec() -> SkillSpec {
SkillSpec {
name: "manage_scheduled_tasks".into(),
description: "管理定时任务:列出/添加/更新/删除/启用禁用。可以管理模型任务(task_type=model),系统任务(task_type=system)只能查看不能修改。需要 PostgreSQL 数据库。".into(),
risk_level: RiskLevel::Low,
parameters: serde_json::json!({
"type": "object",
"properties": {
"action": {
"type": "string",
"enum": ["list", "add", "update", "delete", "toggle"],
"description": "操作类型"
},
"name": {
"type": "string",
"description": "任务名称(add/update 时使用)"
},
"user": {
"type": "string",
"description": "所属用户 IDadd 时使用)"
},
"command": {
"type": "string",
"description": "要执行的 shell 命令(add/update 时使用)"
},
"interval": {
"type": "integer",
"description": "执行间隔(秒,add/update 时使用)"
},
"description": {
"type": "string",
"description": "任务描述(add/update 时使用,可选)"
},
"shell": {
"type": "string",
"description": "Shell 解释器,默认 /bin/bashadd/update 时使用,可选)"
},
"cwd": {
"type": "string",
"description": "工作目录(add/update 时使用,可选)"
},
"id": {
"type": "string",
"description": "任务 UUIDupdate/delete/toggle 时使用)"
},
"task_type": {
"type": "string",
"enum": ["model", "system"],
"description": "任务类型:model(模型任务,默认)或 system(系统任务,仅 add 时使用,可选)"
}
},
"required": ["action"]
}),
timeout_secs: 10,
}
}
/// 执行定时任务管理操作
pub async fn execute(args_json: &str) -> SkillResult {
let pool = match POOL.as_ref() {
Some(p) => p,
None => return SkillResult::error("数据库未配置,请设置 DATABASE_URL 环境变量"),
};
let params: serde_json::Value = serde_json::from_str(args_json).unwrap_or_default();
let action = params
.get("action")
.and_then(|v| v.as_str())
.unwrap_or("list");
match action {
"list" => cmd_list(pool).await,
"add" => cmd_add(pool, &params).await,
"update" => cmd_update(pool, &params).await,
"delete" => cmd_delete(pool, &params).await,
"toggle" => cmd_toggle(pool, &params).await,
_ => SkillResult::error(format!("未知操作: {},支持: list, add, update, delete, toggle", action)),
}
}
/// 检查任务是否为 system 类型(不可修改)
async fn check_system_task(pool: &sqlx::PgPool, id: &uuid::Uuid) -> Result<bool, String> {
let row = sqlx::query_scalar::<_, String>(
"SELECT task_type FROM scheduled_tasks WHERE id = $1",
)
.bind(id)
.fetch_optional(pool)
.await
.map_err(|e| format!("查询任务类型失败: {}", e))?;
match row {
Some(t) => Ok(t == "system"),
None => Err("未找到该任务".to_string()),
}
}
async fn cmd_list(pool: &sqlx::PgPool) -> SkillResult {
let tasks = match crate::db::models::list_scheduled_tasks(pool).await {
Ok(t) => t,
Err(e) => return SkillResult::error(format!("查询失败: {}", e)),
};
if tasks.is_empty() {
return SkillResult::ok("暂无定时任务。你可以通过「添加定时任务」来创建,例如每天早8点播报天气。");
}
let mut lines = vec![format!("📋 共有 {} 个定时任务:\n", tasks.len())];
for task in &tasks {
let status = if task.enabled { "🟢" } else { "🔴" };
let type_tag = if task.task_type == "system" { " [系统]" } else { "" };
let interval = if task.interval_seconds >= 86400 {
format!("{}", task.interval_seconds / 86400)
} else if task.interval_seconds >= 3600 {
format!("{}小时", task.interval_seconds / 3600)
} else {
format!("{}", task.interval_seconds)
};
let last = task
.last_run_at
.map(|t| t.format("%m-%d %H:%M").to_string())
.unwrap_or_else(|| "-".to_string());
let next = task.next_run_at.format("%m-%d %H:%M");
lines.push(format!(
"{status}{type_tag} {name}\n ID: {id}\n 命令: {cmd}\n 间隔: {interval}\n 上次: {last} 下次: {next}",
status = status,
type_tag = type_tag,
name = task.name,
id = task.id,
cmd = task.command,
interval = interval,
last = last,
next = next,
));
}
SkillResult::ok(lines.join("\n"))
}
async fn cmd_add(pool: &sqlx::PgPool, params: &serde_json::Value) -> SkillResult {
let name = params.get("name").and_then(|v| v.as_str()).unwrap_or("");
let user = params.get("user").and_then(|v| v.as_str()).unwrap_or("");
let command = params.get("command").and_then(|v| v.as_str()).unwrap_or("");
let interval = params.get("interval").and_then(|v| v.as_i64()).unwrap_or(0) as i32;
let task_type = params.get("task_type").and_then(|v| v.as_str()).unwrap_or("model");
if name.is_empty() {
return SkillResult::error("请提供 name(任务名称)");
}
if user.is_empty() {
return SkillResult::error("请提供 user(所属用户 ID");
}
if command.is_empty() {
return SkillResult::error("请提供 command(要执行的命令)");
}
if interval <= 0 {
return SkillResult::error("请提供 interval(执行间隔,大于0秒)");
}
if task_type != "model" && task_type != "system" {
return SkillResult::error("task_type 必须是 model 或 system");
}
let description = params.get("description").and_then(|v| v.as_str()).unwrap_or("");
let shell = params.get("shell").and_then(|v| v.as_str()).unwrap_or("/bin/bash");
let cwd = params.get("cwd").and_then(|v| v.as_str()).unwrap_or("");
match crate::db::models::add_scheduled_task(pool, name, user, command, interval, description, shell, cwd, task_type).await {
Ok(id) => SkillResult::ok(format!(
"✅ 定时任务已添加\n 名称: {name}\n 命令: {command}\n 间隔: {interval}\n 类型: {task_type}\n ID: {id}\n\n任务将按设定间隔自动执行,执行结果会通知你。",
name = name, command = command, interval = interval, task_type = task_type, id = id
)),
Err(e) => SkillResult::error(format!("添加失败: {}", e)),
}
}
async fn cmd_update(pool: &sqlx::PgPool, params: &serde_json::Value) -> SkillResult {
let id_str = params.get("id").and_then(|v| v.as_str()).unwrap_or("");
let id = match uuid::Uuid::parse_str(id_str) {
Ok(u) => u,
Err(_) => return SkillResult::error("请提供有效的 id(任务 UUID"),
};
// 检查是否为系统任务
match check_system_task(pool, &id).await {
Ok(true) => return SkillResult::error("❌ 系统任务不可修改。系统任务由管理员配置,如需变更请联系管理员。"),
Ok(false) => {} // model 任务,可以修改
Err(e) => return SkillResult::error(e),
}
let name = params.get("name").and_then(|v| v.as_str());
let command = params.get("command").and_then(|v| v.as_str());
let interval = params.get("interval").and_then(|v| v.as_i64()).map(|v| v as i32);
let description = params.get("description").and_then(|v| v.as_str());
let shell = params.get("shell").and_then(|v| v.as_str());
let cwd = params.get("cwd").and_then(|v| v.as_str());
if name.is_none() && command.is_none() && interval.is_none()
&& description.is_none() && shell.is_none() && cwd.is_none()
{
return SkillResult::error("请至少提供要更新的字段(name/command/interval/description/shell/cwd");
}
match crate::db::models::update_scheduled_task(
pool, &id, name, command, interval, description, shell, cwd,
)
.await
{
Ok(true) => SkillResult::ok("✅ 定时任务已更新"),
Ok(false) => SkillResult::error("未找到该任务或无字段更新"),
Err(e) => SkillResult::error(format!("更新失败: {}", e)),
}
}
async fn cmd_delete(pool: &sqlx::PgPool, params: &serde_json::Value) -> SkillResult {
let id_str = params.get("id").and_then(|v| v.as_str()).unwrap_or("");
let id = match uuid::Uuid::parse_str(id_str) {
Ok(u) => u,
Err(_) => return SkillResult::error("请提供有效的 id(任务 UUID"),
};
// 检查是否为系统任务
match check_system_task(pool, &id).await {
Ok(true) => return SkillResult::error("❌ 系统任务不可删除。系统任务由管理员配置,如需变更请联系管理员。"),
Ok(false) => {} // model 任务,可以删除
Err(e) => return SkillResult::error(e),
}
match crate::db::models::delete_scheduled_task(pool, &id).await {
Ok(true) => SkillResult::ok("✅ 定时任务已删除"),
Ok(false) => SkillResult::error("未找到该任务"),
Err(e) => SkillResult::error(format!("删除失败: {}", e)),
}
}
async fn cmd_toggle(pool: &sqlx::PgPool, params: &serde_json::Value) -> SkillResult {
let id_str = params.get("id").and_then(|v| v.as_str()).unwrap_or("");
let id = match uuid::Uuid::parse_str(id_str) {
Ok(u) => u,
Err(_) => return SkillResult::error("请提供有效的 id(任务 UUID"),
};
// 检查是否为系统任务
match check_system_task(pool, &id).await {
Ok(true) => return SkillResult::error("❌ 系统任务不可启用/禁用。系统任务由管理员配置,如需变更请联系管理员。"),
Ok(false) => {} // model 任务,可以切换
Err(e) => return SkillResult::error(e),
}
match crate::db::models::toggle_scheduled_task(pool, &id).await {
Ok((true, enabled)) => {
let status = if enabled { "🟢 已启用" } else { "🔴 已禁用" };
SkillResult::ok(format!("✅ 定时任务状态已切换: {}", status))
}
Ok((false, _)) => SkillResult::error("未找到该任务"),
Err(e) => SkillResult::error(format!("切换失败: {}", e)),
}
}
+164 -62
View File
@@ -1,6 +1,12 @@
//! ## 工具系统 —— 工具 + 分离式内置工具
//! ## 工具系统 —— 抽象工具 + 内置工具 + 审批管理
//!
//! 本模块是 iAs 的"工具系统",让 LLM 能够调用外部功能。
//! 所有工具分为三种类型:
//!
//! | 类型 | 说明 | 示例 |
//! |------|------|------|
//! | `ApiTool` | Rust 代码实现的工具(HTTP API / 本地计算 / 文件 I/O) | 天气、搜索、高德地图 |
//! | `StdioTool` | 通过子进程 stdin/stdout 通信的工具 | 外部脚本、可执行程序 |
//! | `ShellTool` | 通过 shell 命令执行的工具 | 定时任务、系统命令 |
//!
//! ### 架构设计(两层元工具)
//!
@@ -9,51 +15,42 @@
//! 1. **`query_capabilities`** — 列出所有可用工具的名称、描述、参数格式
//! 2. **`call_capability`** — 按名称调用工具,参数通过 JSON 传递
//!
//! 这样做的优势:新增工具时只需要在服务端注册,不需要修改 LLM 的 tool definitions。
//!
//! ### 工具类型
//! * **内置工具** (builtins/) — Rust 代码实现的工具(天气、搜索、高德地图等)
//! * **上下文工具** — read_memories, write_memory, read_summaries(直接在 daemon 中处理)
//!
//! ### 工具执行流程
//! ```text
//! LLM 请求 call_capability("query_weather", {"location":"北京"})
//! → daemon 解析出 target_name="query_weather",提取参数
//! → 高风险?→ 走 ApprovalManager 审批流
//! → 低风险?→ BuiltinRegistry::execute() 或 SubprocessRunner::execute()
//! → 结果返回 LLM 继续对话
//! ## 工具系统 —— 元工具 + 内置工具 + 审批管理
//!
//! 本模块是 iAs 的"工具系统",让 LLM 能够调用外部功能。
//!
//! ### 架构设计(两层元工具)
//!
//! LLM 不直接感知每个具体工具,而是通过两个**元工具**间接调用:
//!
//! 1. **`query_capabilities`** — 列出所有可用工具的名称、描述、参数格式
//! 2. **`call_capability`** — 按名称调用工具,参数通过 JSON 传递
//!
//! 这样做的优势:新增工具时只需要在服务端注册,不需要修改 LLM 的 tool definitions。
//!
//! ### 工具类型
//! - **内置工具** (builtins/) — Rust 代码实现的工具(天气、搜索、高德地图等)
//! - **上下文工具** — read_memories, write_memory, read_summaries(直接在 daemon 中处理)
//!
//! ### 工具执行流程
//! ```text
//! LLM 请求 call_capability("query_weather", {"location":"北京"})
//! → daemon 解析出 target_name="query_weather",提取参数
//! → 高风险?→ 走 ApprovalManager 审批流
//! → 低风险?→ BuiltinRegistry::execute()
//! → 低风险?→ ToolRegistry::execute()
//! → 结果返回 LLM 继续对话
//! ```
pub mod api_tool;
pub mod approval;
pub mod builtin;
pub mod builtins;
pub mod shell_tool;
pub mod stdio_tool;
pub mod tool;
pub mod types;
/// 重新导出核心类型
pub use api_tool::ApiTool;
pub use tool::ToolRegistry;
pub use types::{SkillResult, SkillSpec};
// 以下类型已定义但尚未被现有代码直接引用
// 新工具可通过这些类型实现:
// - ApiTool: Rust 代码实现的工具
// - StdioTool: 子进程 stdin/stdout 通信的工具
// - ShellTool: shell 命令执行的工具
// - Tool trait: 所有工具的基类
#[allow(unused_imports)]
pub use shell_tool::ShellTool;
#[allow(unused_imports)]
pub use stdio_tool::StdioTool;
#[allow(unused_imports)]
pub use tool::Tool;
/// ## 从 `call_capability` 的 `{name, prompt}` JSON 中提取目标工具的真实参数
///
/// `call_capability` 的参数格式为 `{name: "工具名", prompt: "参数JSON字符串"}`。
@@ -95,40 +92,112 @@ pub fn unpack_call_params(cp: &serde_json::Value) -> serde_json::Value {
serde_json::Value::Object(params)
}
#[cfg(test)]
mod tests {
#[test]
fn test_unpack_call_params() {
use super::unpack_call_params;
use std::sync::LazyLock;
// prompt 是 JSON 字符串
let cp = serde_json::json!({"name": "amap_poi_search", "prompt": "{\"keywords\":\"肯德基\",\"city\":\"北京\"}"});
let result = unpack_call_params(&cp);
assert_eq!(result["keywords"], "肯德基");
assert_eq!(result["city"], "北京");
assert!(result.get("name").is_none());
assert!(result.get("prompt").is_none());
// ═══════════════════════════════════════════════
// 全局工具注册表(主入口)
// ═══════════════════════════════════════════════
// prompt 是 JSON 对象
let cp = serde_json::json!({"name": "query_weather", "prompt": {"location": "上海", "days": 7}});
let result = unpack_call_params(&cp);
assert_eq!(result["location"], "上海");
assert_eq!(result["days"], 7);
/// 全局内置工具注册表
pub static REGISTRY: LazyLock<ToolRegistry> = LazyLock::new(build_registry);
// 仅有 name,无 prompt
let cp = serde_json::json!({"name": "get_current_datetime"});
let result = unpack_call_params(&cp);
assert!(result.as_object().unwrap().is_empty());
/// 构建注册表:注册所有内置工具
fn build_registry() -> ToolRegistry {
let mut registry = ToolRegistry::new();
// 顶层字段合并
let cp = serde_json::json!({"name": "call_capability", "prompt": "{\"a\": 1}", "extra": 42});
let result = unpack_call_params(&cp);
assert_eq!(result["a"], 1);
assert_eq!(result["extra"], 42);
// ── ApiTool: 日期时间(纯本地计算) ──
registry.register(Box::new(ApiTool::new(
self::builtins::datetime::spec(),
std::sync::Arc::new(|_params| {
Box::pin(async { self::builtins::datetime::execute() })
}),
)));
// ── ApiTool: 备忘录(文件 I/O ──
let memos_spec = self::builtins::memos::spec();
registry.register(Box::new(ApiTool::new(
memos_spec,
std::sync::Arc::new(|params| {
Box::pin(async move {
let args = serde_json::to_string(&params).unwrap_or_default();
self::builtins::memos::execute(&args).await
})
}),
)));
// ── ApiTool: 天气(HTTP API ──
registry.register(Box::new(ApiTool::new(
self::builtins::weather::spec(),
std::sync::Arc::new(|params| {
Box::pin(async { self::builtins::weather::execute(params).await })
}),
)));
// ── ApiTool: 联网搜索(HTTP API ──
registry.register(Box::new(ApiTool::new(
self::builtins::web_search::spec(),
std::sync::Arc::new(|params| {
Box::pin(async { self::builtins::web_search::execute(params).await })
}),
)));
// ── ApiTool: 网页抓取(HTTP API ──
registry.register(Box::new(ApiTool::new(
self::builtins::fetch_page::spec(),
std::sync::Arc::new(|params| {
Box::pin(async { self::builtins::fetch_page::execute(params).await })
}),
)));
// ── ApiTool: 高德地图系列(HTTP API ──
for amap_spec in self::builtins::amap::specs() {
let name = amap_spec.name.clone();
registry.register(Box::new(ApiTool::new(
amap_spec,
std::sync::Arc::new(move |params| {
let n = name.clone();
Box::pin(async move {
self::builtins::amap::execute(&n, params)
.await
.unwrap_or_else(|| SkillResult::error("执行失败"))
})
}),
)));
}
// ── ApiTool: 定时任务管理(PostgreSQL ──
registry.register(Box::new(ApiTool::new(
self::builtins::scheduled_task::spec(),
std::sync::Arc::new(|params| {
Box::pin(async move {
let args = serde_json::to_string(&params).unwrap_or_default();
self::builtins::scheduled_task::execute(&args).await
})
}),
)));
registry
}
use types::SkillSpec;
/// 按名称执行工具(主入口)
pub async fn execute(name: &str, args_json: &str) -> Option<SkillResult> {
REGISTRY.execute(name, args_json).await
}
/// 返回所有工具的 spec 列表
pub fn specs() -> Vec<SkillSpec> {
REGISTRY.specs().into_iter().cloned().collect()
}
/// 判断工具是否为高风险
pub fn is_high_risk(name: &str) -> bool {
REGISTRY.is_high_risk(name)
}
/// 判断名称是否对应一个已注册的工具
pub fn is_builtin(name: &str) -> bool {
REGISTRY.is_builtin(name)
}
/// 构建工具调用指南(含优先级说明)
/// 返回给 LLM 的工具列表 + 使用建议
@@ -157,7 +226,7 @@ pub fn build_capability_guide(specs: &[SkillSpec]) -> String {
format_spec(&mut lines, spec);
}
lines.push(String::new());
lines.push(" ▸ 任何天气相关问题(实时天气、天气预报、逐时预报)→ 使用 query_weather".to_string());
lines.push(" ▸ 任何天气相关问题(实时天气、预报、逐时预报)→ 使用 query_weather".to_string());
lines.push(String::new());
lines.push("## 🔍 其他工具:".to_string());
@@ -199,3 +268,36 @@ fn format_spec(lines: &mut Vec<String>, spec: &SkillSpec) {
));
}
}
#[cfg(test)]
mod tests {
#[test]
fn test_unpack_call_params() {
use super::unpack_call_params;
// prompt 是 JSON 字符串
let cp = serde_json::json!({"name": "amap_poi_search", "prompt": "{\"keywords\":\"肯德基\",\"city\":\"北京\"}"});
let result = unpack_call_params(&cp);
assert_eq!(result["keywords"], "肯德基");
assert_eq!(result["city"], "北京");
assert!(result.get("name").is_none());
assert!(result.get("prompt").is_none());
// prompt 是 JSON 对象
let cp = serde_json::json!({"name": "query_weather", "prompt": {"location": "上海", "days": 7}});
let result = unpack_call_params(&cp);
assert_eq!(result["location"], "上海");
assert_eq!(result["days"], 7);
// 仅有 name,无 prompt
let cp = serde_json::json!({"name": "get_current_datetime"});
let result = unpack_call_params(&cp);
assert!(result.as_object().unwrap().is_empty());
// 顶层字段合并
let cp = serde_json::json!({"name": "call_capability", "prompt": "{\"a\": 1}", "extra": 42});
let result = unpack_call_params(&cp);
assert_eq!(result["a"], 1);
assert_eq!(result["extra"], 42);
}
}
+117
View File
@@ -0,0 +1,117 @@
//! ## Shell 工具 —— 通过 shell 命令执行的工具
//!
//! 执行一条 shell 命令(通过 `/bin/bash -c` 或自定义 shell),
//! 将参数通过环境变量 `TOOL_PARAMS` 传入,从 stdout 获取结果。
//!
//! ### 参数传递
//! 参数以 JSON 字符串形式通过环境变量 `TOOL_PARAMS` 传入子进程。
//!
//! ### 输出格式
//! stdout 的内容作为工具的输出结果返回。
//! 如果输出是 JSON 格式的 `SkillResult`,则解析后返回结构化结果。
use async_trait::async_trait;
use std::time::Duration;
use tokio::process::Command;
use super::tool::Tool;
use super::types::{SkillResult, SkillSpec};
/// ## Shell 工具
///
/// 通过 shell 命令执行的工具。
///
/// 当前未被现有代码直接引用,供后续扩展使用。
#[allow(dead_code)]
pub struct ShellTool {
spec: SkillSpec,
command: String,
shell: String,
cwd: Option<String>,
timeout: Duration,
}
#[allow(dead_code)]
impl ShellTool {
/// 创建 Shell 工具
///
/// ### 参数
/// - `spec` — 工具元数据
/// - `command` — shell 命令字符串
/// - `timeout_secs` — 超时秒数(默认 30
pub fn new(spec: SkillSpec, command: impl Into<String>, timeout_secs: u64) -> Self {
Self {
spec,
command: command.into(),
shell: "/bin/bash".into(),
cwd: None,
timeout: Duration::from_secs(timeout_secs.max(5)),
}
}
/// 设置 shell 解释器
pub fn with_shell(mut self, shell: impl Into<String>) -> Self {
self.shell = shell.into();
self
}
/// 设置工作目录
pub fn with_cwd(mut self, cwd: impl Into<String>) -> Self {
self.cwd = Some(cwd.into());
self
}
}
#[async_trait]
impl Tool for ShellTool {
fn spec(&self) -> &SkillSpec {
&self.spec
}
async fn execute(&self, params: serde_json::Value) -> SkillResult {
let params_str = serde_json::to_string(&params).unwrap_or_default();
let mut cmd = Command::new(&self.shell);
cmd.arg("-c")
.arg(&self.command)
.env("TOOL_PARAMS", &params_str)
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.kill_on_drop(true);
if let Some(ref cwd) = self.cwd {
cmd.current_dir(cwd);
}
let output = match tokio::time::timeout(self.timeout, cmd.output()).await {
Ok(Ok(output)) => output,
Ok(Err(e)) => return SkillResult::error(format!("子进程错误: {}", e)),
Err(_) => {
return SkillResult::error(format!("执行超时 ({}s)", self.timeout.as_secs()))
}
};
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
return SkillResult::error(format!(
"退出码 {}: {}",
output.status.code().unwrap_or(-1),
if stderr.is_empty() { stdout } else { stderr }
));
}
let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
if stdout.is_empty() {
return SkillResult::ok("执行成功(无输出)");
}
// 尝试解析 JSON 结果
if let Ok(result) = serde_json::from_str::<SkillResult>(&stdout) {
return result;
}
// 非 JSON 输出,当作纯文本结果
SkillResult::ok(stdout)
}
}
+116
View File
@@ -0,0 +1,116 @@
//! ## Stdio 工具 —— 通过子进程 stdin/stdout 通信的工具
//!
//! 启动一个子进程,通过 stdin 发送 JSON 参数,从 stdout 读取 JSON 结果。
//! 适用于外部脚本、可执行程序等。
//!
//! ### 协议
//! - stdin: 一行 JSON 字符串
//! - stdout: 一行 JSON 字符串(`{"success": true/false, "output": "...", "data": ...}`
//! - 超时: 默认 30 秒,可通过 `timeout_secs` 配置
use async_trait::async_trait;
use std::time::Duration;
use tokio::process::Command;
use tokio::io::AsyncWriteExt;
use super::tool::Tool;
use super::types::{SkillResult, SkillSpec};
/// ## Stdio 工具
///
/// 通过子进程 stdin/stdout 通信的工具。
///
/// 当前未被现有代码直接引用,供后续扩展使用。
#[allow(dead_code)]
pub struct StdioTool {
spec: SkillSpec,
command: String,
args: Vec<String>,
timeout: Duration,
}
impl StdioTool {
/// 创建 Stdio 工具
#[allow(dead_code)]
///
/// ### 参数
/// - `spec` — 工具元数据
/// - `command` — 可执行文件路径
/// - `args` — 启动参数
/// - `timeout_secs` — 超时秒数(默认 30
pub fn new(
spec: SkillSpec,
command: impl Into<String>,
args: Vec<String>,
timeout_secs: u64,
) -> Self {
Self {
spec,
command: command.into(),
args,
timeout: Duration::from_secs(timeout_secs.max(5)),
}
}
}
#[async_trait]
impl Tool for StdioTool {
fn spec(&self) -> &SkillSpec {
&self.spec
}
async fn execute(&self, params: serde_json::Value) -> SkillResult {
let input = serde_json::to_string(&params).unwrap_or_default();
let mut child = match Command::new(&self.command)
.args(&self.args)
.stdin(std::process::Stdio::piped())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.kill_on_drop(true)
.spawn()
{
Ok(c) => c,
Err(e) => return SkillResult::error(format!("启动子进程失败: {}", e)),
};
// 写入 stdin
if let Some(mut stdin) = child.stdin.take() {
if let Err(e) = stdin.write_all(input.as_bytes()).await {
return SkillResult::error(format!("写入 stdin 失败: {}", e));
}
// 关闭 stdin 表示输入结束
let _ = stdin.shutdown().await;
}
// 等待子进程完成(带超时)
let output = match tokio::time::timeout(self.timeout, child.wait_with_output()).await {
Ok(Ok(output)) => output,
Ok(Err(e)) => return SkillResult::error(format!("子进程错误: {}", e)),
Err(_) => return SkillResult::error(format!("执行超时 ({}s)", self.timeout.as_secs())),
};
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
return SkillResult::error(format!(
"退出码 {}: {}",
output.status.code().unwrap_or(-1),
if stderr.is_empty() { stdout } else { stderr }
));
}
let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
if stdout.is_empty() {
return SkillResult::ok("执行成功(无输出)");
}
// 尝试解析 JSON 结果
if let Ok(result) = serde_json::from_str::<SkillResult>(&stdout) {
return result;
}
// 非 JSON 输出,当作纯文本结果
SkillResult::ok(stdout)
}
}
+91
View File
@@ -0,0 +1,91 @@
//! ## 工具抽象 —— Tool trait + ToolRegistry
//!
//! 定义了所有工具的基类 `Tool` trait,以及统一的注册表 `ToolRegistry`。
//! 所有具体工具(API / Stdio / Shell)都实现此 trait。
use async_trait::async_trait;
use std::collections::HashMap;
use super::types::{RiskLevel, SkillResult, SkillSpec};
/// Boxed future for async tool handlers
pub type BoxFuture<T> = std::pin::Pin<Box<dyn std::future::Future<Output = T> + Send>>;
/// Async handler function: JSON params → SkillResult
pub type ToolHandler =
std::sync::Arc<dyn Fn(serde_json::Value) -> BoxFuture<SkillResult> + Send + Sync>;
/// ## 工具基类
///
/// 所有工具(API / Stdio / Shell)都实现此 trait。
///
/// ### 方法
/// - `spec()` — 返回工具的元数据(名称、描述、参数、风险等级)
/// - `execute()` — 执行工具,接收 JSON 参数,返回执行结果
#[async_trait]
pub trait Tool: Send + Sync {
fn spec(&self) -> &SkillSpec;
async fn execute(&self, params: serde_json::Value) -> SkillResult;
}
/// ## 工具注册表
///
/// 管理所有已注册的工具,提供按名称查找和执行的能力。
///
/// ### 用法
/// ```rust,ignore
/// let mut registry = ToolRegistry::new();
/// registry.register(Box::new(api_tool));
/// registry.register(Box::new(shell_tool));
///
/// let result = registry.execute("tool_name", r#"{"key":"value"}"#).await;
/// ```
pub struct ToolRegistry {
tools: HashMap<String, Box<dyn Tool>>,
}
impl ToolRegistry {
pub fn new() -> Self {
Self {
tools: HashMap::new(),
}
}
/// 注册一个工具
pub fn register(&mut self, tool: Box<dyn Tool>) {
let name = tool.spec().name.clone();
self.tools.insert(name, tool);
}
/// 按名称执行工具
pub async fn execute(&self, name: &str, params_json: &str) -> Option<SkillResult> {
let params: serde_json::Value = serde_json::from_str(params_json).unwrap_or_default();
match self.tools.get(name) {
Some(tool) => Some(tool.execute(params).await),
None => None,
}
}
/// 按名称获取工具引用
#[allow(dead_code)]
pub fn get(&self, name: &str) -> Option<&dyn Tool> {
self.tools.get(name).map(|t| t.as_ref())
}
/// 返回所有工具的 spec 引用列表
pub fn specs(&self) -> Vec<&SkillSpec> {
self.tools.values().map(|t| t.spec()).collect()
}
/// 判断名称是否对应一个已注册的工具
pub fn is_builtin(&self, name: &str) -> bool {
self.tools.contains_key(name)
}
/// 判断工具是否为高风险(需要审批)
pub fn is_high_risk(&self, name: &str) -> bool {
self.tools
.get(name)
.map_or(false, |t| t.spec().risk_level == RiskLevel::High)
}
}
+1 -1
View File
@@ -62,7 +62,7 @@ fn default_timeout() -> u64 {
/// * `output` — 可读的输出文本(将返回给 LLM)
/// * `data` — 结构化数据(可选),用于后续处理
#[derive(Debug, Clone, Serialize)]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SkillResult {
pub success: bool,
pub output: String,
+3 -3
View File
@@ -234,7 +234,7 @@ fn build_worker_executor(
}
}
"query_capabilities" => {
let specs = crate::tools::builtin::BuiltinRegistry::specs();
let specs = crate::tools::specs();
Ok(crate::tools::build_capability_guide(&specs))
}
// call_capability 或直接调用
@@ -257,7 +257,7 @@ fn build_worker_executor(
}
// 高风险工具 → 请求审批(除非已被用户批准)
if crate::tools::builtin::BuiltinRegistry::is_high_risk(&target_name) {
if crate::tools::is_high_risk(&target_name) {
let is_approved = shared.approved_tool.lock().await.as_deref() == Some(&target_name);
if !is_approved {
let reason = format!("LLM 请求执行高风险工具: {}", target_name);
@@ -270,7 +270,7 @@ fn build_worker_executor(
// 执行内置工具
if let Some(result) =
crate::tools::builtin::BuiltinRegistry::execute(&target_name, &target_args)
crate::tools::execute(&target_name, &target_args)
.await
{
if target_name == "manage_memos" {