diff --git a/Cargo.toml b/Cargo.toml index 1b85fcd..c45d272 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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" diff --git a/migrations/20260601000014_scheduled_tasks_type.sql b/migrations/20260601000014_scheduled_tasks_type.sql new file mode 100644 index 0000000..6ee32af --- /dev/null +++ b/migrations/20260601000014_scheduled_tasks_type.sql @@ -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); diff --git a/src/bin/ias-auth.rs b/src/bin/ias-auth.rs deleted file mode 100644 index 538224d..0000000 --- a/src/bin/ias-auth.rs +++ /dev/null @@ -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 = std::env::args().collect(); - if args.len() < 2 { - eprintln!("用法: ias-auth "); - 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 { - 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::>() - .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()) - )) -} diff --git a/src/cli.rs b/src/cli.rs index b754a85..8973ffe 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -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 --name "新名称" --interval 3600 + /// ias task delete + /// ias task toggle + #[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, + + /// Shell 解释器(默认 /bin/bash) + #[arg(long)] + shell: Option, + + /// 工作目录 + #[arg(long)] + cwd: Option, + }, + + /// 更新定时任务 + /// + /// 示例: + /// ias task update --name "新名称" --interval 3600 + Update { + /// 任务 ID(UUID) + id: String, + + /// 新名称 + #[arg(long)] + name: Option, + + /// 新命令 + #[arg(long)] + command: Option, + + /// 新执行间隔(秒) + #[arg(long)] + interval: Option, + + /// 新描述 + #[arg(long)] + description: Option, + + /// 新 Shell 解释器 + #[arg(long)] + shell: Option, + + /// 新工作目录 + #[arg(long)] + cwd: Option, + }, + + /// 删除定时任务 + /// + /// 示例: + /// ias task delete + Delete { + /// 任务 ID(UUID) + id: String, + }, + + /// 启用/禁用定时任务 + /// + /// 示例: + /// ias task toggle + Toggle { + /// 任务 ID(UUID) + id: String, + }, +} + // ─── 工具子命令 ─── #[derive(Subcommand, Debug, Clone)] diff --git a/src/daemon.rs b/src/daemon.rs index e02d0d7..dd0178f 100644 --- a/src/daemon.rs +++ b/src/daemon.rs @@ -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::>().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::>().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 { } })); - let specs = crate::tools::builtin::BuiltinRegistry::specs(); + let specs = crate::tools::specs(); for spec in &specs { list.push(serde_json::json!({ "type": "function", diff --git a/src/db/models.rs b/src/db/models.rs index 37a4346..ef11e6b 100644 --- a/src/db/models.rs +++ b/src/db/models.rs @@ -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, + pub command: String, + pub shell: String, + pub cwd: String, + pub interval_seconds: i32, + pub enabled: bool, + pub next_run_at: DateTime, + pub last_run_at: Option>, + pub last_status: Option, + pub task_type: String, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +/// 列出所有定时任务 +pub async fn list_scheduled_tasks(pool: &PgPool) -> Result, 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 { + 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, + description: Option<&str>, + shell: Option<&str>, + cwd: Option<&str>, +) -> Result { + 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 { + 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)), + } +} diff --git a/src/main.rs b/src/main.rs index 56e6db2..050e36c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -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>) { + 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) { diff --git a/src/tools/api_tool.rs b/src/tools/api_tool.rs new file mode 100644 index 0000000..1feb9d7 --- /dev/null +++ b/src/tools/api_tool.rs @@ -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 + } +} diff --git a/src/tools/builtin.rs b/src/tools/builtin.rs index b30cc3e..a1e8d0a 100644 --- a/src/tools/builtin.rs +++ b/src/tools/builtin.rs @@ -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 { - 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 { - 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 { + 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) } } diff --git a/src/tools/builtins/mod.rs b/src/tools/builtins/mod.rs index eb53c9d..722585f 100644 --- a/src/tools/builtins/mod.rs +++ b/src/tools/builtins/mod.rs @@ -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; diff --git a/src/tools/builtins/scheduled_task.rs b/src/tools/builtins/scheduled_task.rs new file mode 100644 index 0000000..4a3af29 --- /dev/null +++ b/src/tools/builtins/scheduled_task.rs @@ -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> = 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": "所属用户 ID(add 时使用)" + }, + "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/bash(add/update 时使用,可选)" + }, + "cwd": { + "type": "string", + "description": "工作目录(add/update 时使用,可选)" + }, + "id": { + "type": "string", + "description": "任务 UUID(update/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, ¶ms).await, + "update" => cmd_update(pool, ¶ms).await, + "delete" => cmd_delete(pool, ¶ms).await, + "toggle" => cmd_toggle(pool, ¶ms).await, + _ => SkillResult::error(format!("未知操作: {},支持: list, add, update, delete, toggle", action)), + } +} + +/// 检查任务是否为 system 类型(不可修改) +async fn check_system_task(pool: &sqlx::PgPool, id: &uuid::Uuid) -> Result { + 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)), + } +} diff --git a/src/tools/mod.rs b/src/tools/mod.rs index d0a78fe..bca4f7c 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -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 = 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(¶ms).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(¶ms).unwrap_or_default(); + self::builtins::scheduled_task::execute(&args).await + }) + }), + ))); + + registry } -use types::SkillSpec; +/// 按名称执行工具(主入口) +pub async fn execute(name: &str, args_json: &str) -> Option { + REGISTRY.execute(name, args_json).await +} + +/// 返回所有工具的 spec 列表 +pub fn specs() -> Vec { + 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, 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); + } +} diff --git a/src/tools/shell_tool.rs b/src/tools/shell_tool.rs new file mode 100644 index 0000000..608caaa --- /dev/null +++ b/src/tools/shell_tool.rs @@ -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, + timeout: Duration, +} + +#[allow(dead_code)] +impl ShellTool { + /// 创建 Shell 工具 + /// + /// ### 参数 + /// - `spec` — 工具元数据 + /// - `command` — shell 命令字符串 + /// - `timeout_secs` — 超时秒数(默认 30) + pub fn new(spec: SkillSpec, command: impl Into, 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) -> Self { + self.shell = shell.into(); + self + } + + /// 设置工作目录 + pub fn with_cwd(mut self, cwd: impl Into) -> 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(¶ms).unwrap_or_default(); + + let mut cmd = Command::new(&self.shell); + cmd.arg("-c") + .arg(&self.command) + .env("TOOL_PARAMS", ¶ms_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::(&stdout) { + return result; + } + + // 非 JSON 输出,当作纯文本结果 + SkillResult::ok(stdout) + } +} diff --git a/src/tools/stdio_tool.rs b/src/tools/stdio_tool.rs new file mode 100644 index 0000000..b0076ab --- /dev/null +++ b/src/tools/stdio_tool.rs @@ -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, + timeout: Duration, +} + +impl StdioTool { + /// 创建 Stdio 工具 + #[allow(dead_code)] + /// + /// ### 参数 + /// - `spec` — 工具元数据 + /// - `command` — 可执行文件路径 + /// - `args` — 启动参数 + /// - `timeout_secs` — 超时秒数(默认 30) + pub fn new( + spec: SkillSpec, + command: impl Into, + args: Vec, + 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(¶ms).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::(&stdout) { + return result; + } + + // 非 JSON 输出,当作纯文本结果 + SkillResult::ok(stdout) + } +} diff --git a/src/tools/tool.rs b/src/tools/tool.rs new file mode 100644 index 0000000..be44c0c --- /dev/null +++ b/src/tools/tool.rs @@ -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 = std::pin::Pin + Send>>; + +/// Async handler function: JSON params → SkillResult +pub type ToolHandler = + std::sync::Arc BoxFuture + 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>, +} + +impl ToolRegistry { + pub fn new() -> Self { + Self { + tools: HashMap::new(), + } + } + + /// 注册一个工具 + pub fn register(&mut self, tool: Box) { + let name = tool.spec().name.clone(); + self.tools.insert(name, tool); + } + + /// 按名称执行工具 + pub async fn execute(&self, name: &str, params_json: &str) -> Option { + 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) + } +} diff --git a/src/tools/types.rs b/src/tools/types.rs index 6bffb5c..2200eb7 100644 --- a/src/tools/types.rs +++ b/src/tools/types.rs @@ -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, diff --git a/src/worker.rs b/src/worker.rs index f01433a..a9b68f0 100644 --- a/src/worker.rs +++ b/src/worker.rs @@ -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" {