diff --git a/src/main.rs b/src/main.rs index c7a68cc..6e29750 100644 --- a/src/main.rs +++ b/src/main.rs @@ -15,7 +15,7 @@ use context::MemoryStore; use db::Database; use llm::{Conversation, ConversationConfig, DEFAULT_SYSTEM_PROMPT, ToolExecutor, Usage}; use std::sync::Arc; -use tools::approval::ApprovalManager; +use tools::approval::{ApprovalDecision, ApprovalManager}; use tools::registry::{ExecutionContext, SkillRegistry}; use tracing::{error, info}; use wechat::client::WeChatClient; @@ -60,9 +60,6 @@ async fn main() { )); memory_store.load("").await; - // 文件状态管理器(降级方案) - let file_state = state::StateManager::new(".data/weixin-ilink"); - match cli.command { Commands::Login { timeout } => { cmd_login(timeout, &database, &file_state).await; @@ -343,59 +340,63 @@ async fn cmd_listen( Box::pin(async move { let user_id = session.lock().await.current_user_id.clone(); - // 内置工具优先 - if let Some(result) = tools::builtin::BuiltinRegistry::execute(&n, &args).await { + // 上下文工具直接处理 + if n == "read_memories" { return Ok(memory_store.read_for(&user_id).await); } + if n == "write_memory" { + let params: serde_json::Value = serde_json::from_str(&args).unwrap_or_default(); + let content = params.get("content").and_then(|v| v.as_str()).unwrap_or(""); + if content.is_empty() { return Ok("请提供 content 参数".to_string()); } + return Ok(memory_store.write_for(&user_id, content).await); + } + if n == "read_summaries" { return Ok(context::tools::read_summaries(&session).await); } + + // query_capabilities: 合并内置 + 外部(去重) + if n == "query_capabilities" { + let spec_list = if let Some(ref r) = reg { + let mut s: Vec = r.list_specs().into_iter().cloned().collect(); + for builtin in tools::builtin::BuiltinRegistry::specs() { + s.retain(|x| x.name != builtin.name); + s.push(builtin); + } + s + } else { + tools::builtin::BuiltinRegistry::specs() + }; + return Ok(build_capability_list(&spec_list)); + } + + // 确定目标工具名(call_capability 需要提取 name 参数) + let target_name = if n == "call_capability" { + let cp: serde_json::Value = serde_json::from_str(&args).unwrap_or_default(); + cp.get("name").and_then(|v| v.as_str()).map(|s| s.to_string()).unwrap_or_default() + } else { n.clone() }; + + // 内置工具 + if let Some(mut result) = tools::builtin::BuiltinRegistry::execute(&target_name, &args).await { + if tools::builtin::BuiltinRegistry::is_high_risk(&target_name) { + let mut ctx = ExecutionContext::new(&user_id); + ctx = ctx.with_approval(approval.clone()); + ctx = ctx.with_wechat(send); + let approved = builtin_approve(&ctx, &target_name).await; + match approved { + Ok(true) => return Ok(result.output), + Ok(false) => return Ok(SkillResult::rejected(&target_name).output), + Err(e) => return Ok(e), + } + } return Ok(result.output); } - match n.as_str() { - "read_memories" => return Ok(memory_store.read_for(&user_id).await), - "write_memory" => { - let params: serde_json::Value = - serde_json::from_str(&args).unwrap_or_default(); - let content = params.get("content").and_then(|v| v.as_str()).unwrap_or(""); - if content.is_empty() { - return Ok("请提供 content 参数".to_string()); - } - return Ok(memory_store.write_for(&user_id, content).await); - } - "read_summaries" => return Ok(context::tools::read_summaries(&session).await), - "query_capabilities" => { - if let Some(r) = reg.as_ref() { - let mut specs: Vec = r.list_specs().into_iter().cloned().collect(); - specs.extend(tools::builtin::BuiltinRegistry::specs()); - return Ok(build_capability_list(&specs)); - } - return Ok("暂无工具".to_string()); - } - "call_capability" => { - let cap_params: serde_json::Value = serde_json::from_str(&args).unwrap_or_default(); - let cap_name = cap_params.get("name").and_then(|v| v.as_str()).map(|s| s.to_string()); - if let Some(name) = cap_name { - if let Some(ref reg) = reg { - let mut ctx = ExecutionContext::new(&user_id); - ctx = ctx.with_approval(approval.clone()); - ctx = ctx.with_wechat(send); - let result = reg.execute(&name, cap_params, &ctx).await; - return Ok(result.output); - } - } - return Ok("请提供 name 参数(要调用的 skill 名称)".to_string()); - } - _ => {} - } - let user_id = current_user.lock().await.clone(); - if let Some(ref reg) = reg { - let params: serde_json::Value = - serde_json::from_str(&args).unwrap_or(serde_json::json!({})); + // 回退到外部 SkillRegistry + if let Some(ref r) = reg { + let params: serde_json::Value = serde_json::from_str(&args).unwrap_or(serde_json::json!({})); let mut ctx = ExecutionContext::new(&user_id); ctx = ctx.with_approval(approval.clone()); ctx = ctx.with_wechat(send); - let result = reg.execute(&n, params, &ctx).await; - Ok(result.output) - } else { - Ok(format!("未知工具: {}", n)) + let result = r.execute(&target_name, params, &ctx).await; + return Ok(result.output); } + Ok(format!("未知工具: {}", target_name)) }) }); @@ -808,6 +809,29 @@ fn global_skills_dir() -> std::path::PathBuf { .join("skills") } +use tools::skill::SkillResult; + +async fn builtin_approve(ctx: &ExecutionContext, name: &str) -> Result { + let manager = match ctx.approval_manager.as_ref() { + Some(m) => m.clone(), + None => return Ok(true), + }; + let (code, rx) = manager.create(&ctx.user_id, name).await?; + + let msg = format!( + "⚠️ 操作确认\n\n技能:{}\n\n确认码:{}\n\n回复确认码以继续\n回复 0 或「取消」以取消\n(5分钟内有效,最多3次尝试)", + name, code + ); + if let Some(ref send) = ctx.send_wechat { + if let Err(e) = send(&ctx.user_id, &msg).await { return Err(e); } + } + + match tokio::time::timeout(std::time::Duration::from_secs(300), rx).await { + Ok(Ok(ApprovalDecision::Approved)) => Ok(true), + _ => Ok(false), + } +} + fn build_capability_list(specs: &[tools::skill::SkillSpec]) -> String { let mut lines = vec!["📋 可用工具:".to_string()]; for spec in specs { diff --git a/src/tools/builtin.rs b/src/tools/builtin.rs index 6ac149a..232aa27 100644 --- a/src/tools/builtin.rs +++ b/src/tools/builtin.rs @@ -5,7 +5,6 @@ use chrono::Local; pub struct BuiltinRegistry; impl BuiltinRegistry { - /// 执行内置工具,返回 Some(SkillResult) 或 None(不是内置工具) pub async fn execute(name: &str, args_json: &str) -> Option { match name { "get_current_datetime" => Some(execute_datetime()), @@ -33,6 +32,10 @@ impl BuiltinRegistry { }, ] } + + pub fn is_high_risk(name: &str) -> bool { + Self::specs().iter().any(|s| s.name == name && s.risk_level == RiskLevel::High) + } } fn execute_datetime() -> SkillResult { @@ -44,7 +47,7 @@ async fn handle_memos(args_json: &str) -> SkillResult { 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"); let path = ".data/memos.json"; - std::fs::create_dir_all(".data").ok(); + if let Err(e) = std::fs::create_dir_all(".data") { return SkillResult::error(format!("创建目录失败: {}", e)); } let mut data: Vec = if std::path::Path::new(path).exists() { serde_json::from_str(&std::fs::read_to_string(path).unwrap_or_default()).unwrap_or_default()