feat: add tool manager capability
This commit is contained in:
@@ -1,6 +1,7 @@
|
|||||||
# Rust 构建产物
|
# Rust 构建产物
|
||||||
/target
|
/target
|
||||||
tools/*/target/
|
tools/*/target/
|
||||||
|
dist/
|
||||||
|
|
||||||
# 运行时数据(状态文件、缓存、聊天记录本地副本等)
|
# 运行时数据(状态文件、缓存、聊天记录本地副本等)
|
||||||
.data/
|
.data/
|
||||||
|
|||||||
+44
-10
@@ -43,12 +43,13 @@ use crate::context::HistoryEntry;
|
|||||||
use crate::context::MemoryStore;
|
use crate::context::MemoryStore;
|
||||||
use crate::db::Database;
|
use crate::db::Database;
|
||||||
use crate::llm::{
|
use crate::llm::{
|
||||||
ASYNC_MARKER, ChatResult, Conversation, ConversationConfig, DEFAULT_SYSTEM_PROMPT, ToolExecutor,
|
ChatResult, Conversation, ConversationConfig, DEFAULT_SYSTEM_PROMPT, ToolExecution,
|
||||||
|
ToolExecutor,
|
||||||
};
|
};
|
||||||
use crate::queue::{ConsumerChannels, EnqueueHandle, MessageKind, PipelineMessage, QueueRunner};
|
use crate::queue::{ConsumerChannels, EnqueueHandle, MessageKind, PipelineMessage, QueueRunner};
|
||||||
|
|
||||||
use crate::tools::approval::ApprovalManager;
|
use crate::tools::approval::ApprovalManager;
|
||||||
use crate::tools::registry::Registry;
|
use crate::tools::registry::RegistryManager;
|
||||||
use crate::wechat::client::WeChatClient;
|
use crate::wechat::client::WeChatClient;
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
@@ -71,7 +72,7 @@ struct DaemonCtx {
|
|||||||
/// 审批管理器 —— 处理高风险工具的用户确认码流程
|
/// 审批管理器 —— 处理高风险工具的用户确认码流程
|
||||||
approval: Arc<ApprovalManager>,
|
approval: Arc<ApprovalManager>,
|
||||||
/// 工具注册器 —— 管理外部工具规范和执行
|
/// 工具注册器 —— 管理外部工具规范和执行
|
||||||
registry: Arc<Registry>,
|
registry: Arc<RegistryManager>,
|
||||||
/// 长期记忆管理器 —— 按用户隔离的记忆读写
|
/// 长期记忆管理器 —— 按用户隔离的记忆读写
|
||||||
memory_store: Arc<MemoryStore>,
|
memory_store: Arc<MemoryStore>,
|
||||||
/// 系统提示词,覆盖默认值
|
/// 系统提示词,覆盖默认值
|
||||||
@@ -109,7 +110,7 @@ pub async fn run(database: &Arc<Database>, memory_store: &Arc<MemoryStore>) {
|
|||||||
let approval_manager = Arc::new(ApprovalManager::new(Some(Arc::new(
|
let approval_manager = Arc::new(ApprovalManager::new(Some(Arc::new(
|
||||||
database.pool().clone(),
|
database.pool().clone(),
|
||||||
))));
|
))));
|
||||||
let registry = Arc::new(Registry::load("tools"));
|
let registry = Arc::new(RegistryManager::load("tools"));
|
||||||
let system_prompt = std::env::var("WEIXIN_LLM_SYSTEM_PROMPT").unwrap_or_else(|_| String::new());
|
let system_prompt = std::env::var("WEIXIN_LLM_SYSTEM_PROMPT").unwrap_or_else(|_| String::new());
|
||||||
let model = std::env::var("DEEPSEEK_MODEL").unwrap_or_else(|_| "deepseek-v4-flash".to_string());
|
let model = std::env::var("DEEPSEEK_MODEL").unwrap_or_else(|_| "deepseek-v4-flash".to_string());
|
||||||
|
|
||||||
@@ -408,7 +409,7 @@ async fn llm_consumer_loop(
|
|||||||
context_token,
|
context_token,
|
||||||
))
|
))
|
||||||
.await;
|
.await;
|
||||||
Ok(format!("{}{}", ASYNC_MARKER, n))
|
Ok(ToolExecution::Pending)
|
||||||
})
|
})
|
||||||
});
|
});
|
||||||
conv.set_tool_executor(executor);
|
conv.set_tool_executor(executor);
|
||||||
@@ -691,7 +692,8 @@ async fn handle_tool_call(
|
|||||||
eq.enqueue(msg).await;
|
eq.enqueue(msg).await;
|
||||||
})
|
})
|
||||||
};
|
};
|
||||||
ctx.registry
|
let output = ctx
|
||||||
|
.registry
|
||||||
.dispatch(
|
.dispatch(
|
||||||
target,
|
target,
|
||||||
&sub_params,
|
&sub_params,
|
||||||
@@ -701,7 +703,13 @@ async fn handle_tool_call(
|
|||||||
&ctx.memory_store,
|
&ctx.memory_store,
|
||||||
&wechat_cb,
|
&wechat_cb,
|
||||||
)
|
)
|
||||||
.await
|
.await;
|
||||||
|
if target == "tool_manager" && !is_tool_error(&output) {
|
||||||
|
let count = ctx.registry.reload().await;
|
||||||
|
format!("{output}\n\n工具注册表已重载,当前可用工具数: {count}")
|
||||||
|
} else {
|
||||||
|
output
|
||||||
|
}
|
||||||
} else {
|
} else {
|
||||||
let params: serde_json::Value = serde_json::from_str(&arguments).unwrap_or_default();
|
let params: serde_json::Value = serde_json::from_str(&arguments).unwrap_or_default();
|
||||||
let eq = enqueue.clone();
|
let eq = enqueue.clone();
|
||||||
@@ -726,7 +734,8 @@ async fn handle_tool_call(
|
|||||||
eq.enqueue(msg).await;
|
eq.enqueue(msg).await;
|
||||||
})
|
})
|
||||||
};
|
};
|
||||||
ctx.registry
|
let output = ctx
|
||||||
|
.registry
|
||||||
.dispatch(
|
.dispatch(
|
||||||
&tool_name,
|
&tool_name,
|
||||||
¶ms,
|
¶ms,
|
||||||
@@ -736,7 +745,13 @@ async fn handle_tool_call(
|
|||||||
&ctx.memory_store,
|
&ctx.memory_store,
|
||||||
&wechat_cb,
|
&wechat_cb,
|
||||||
)
|
)
|
||||||
.await
|
.await;
|
||||||
|
if tool_name == "tool_manager" && !is_tool_error(&output) {
|
||||||
|
let count = ctx.registry.reload().await;
|
||||||
|
format!("{output}\n\n工具注册表已重载,当前可用工具数: {count}")
|
||||||
|
} else {
|
||||||
|
output
|
||||||
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
enqueue
|
enqueue
|
||||||
@@ -752,6 +767,25 @@ async fn handle_tool_call(
|
|||||||
.await;
|
.await;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn is_tool_error(result: &str) -> bool {
|
||||||
|
result.starts_with("异常退出")
|
||||||
|
|| result.starts_with("启动失败")
|
||||||
|
|| result.starts_with("工具输出非JSON")
|
||||||
|
|| result.starts_with("HTTP失败")
|
||||||
|
|| result.starts_with("读取stdout失败")
|
||||||
|
|| result.starts_with("stdout超时")
|
||||||
|
|| result.starts_with("进程超时")
|
||||||
|
|| result.starts_with("进程异常")
|
||||||
|
|| result.starts_with("未知消息类型")
|
||||||
|
|| result.starts_with("未知工具")
|
||||||
|
|| result.starts_with("不支持的方法")
|
||||||
|
|| result.starts_with("审批失败")
|
||||||
|
|| result.starts_with("审批超时")
|
||||||
|
|| result.starts_with("操作已被用户取消")
|
||||||
|
|| result.starts_with("缺少必填参数")
|
||||||
|
|| result.starts_with("工具 '")
|
||||||
|
}
|
||||||
|
|
||||||
fn unpack_call_capability_params(cp: &serde_json::Value) -> (&str, serde_json::Value) {
|
fn unpack_call_capability_params(cp: &serde_json::Value) -> (&str, serde_json::Value) {
|
||||||
let target = cp["name"].as_str().unwrap_or("");
|
let target = cp["name"].as_str().unwrap_or("");
|
||||||
let mut params = serde_json::Map::new();
|
let mut params = serde_json::Map::new();
|
||||||
@@ -804,7 +838,7 @@ async fn execute_builtin(
|
|||||||
|
|
||||||
match name {
|
match name {
|
||||||
"query_capabilities" => {
|
"query_capabilities" => {
|
||||||
let tools = ctx.registry.list_tools();
|
let tools = ctx.registry.list_tools().await;
|
||||||
if tools.is_empty() {
|
if tools.is_empty() {
|
||||||
return Some("当前没有注册任何工具".into());
|
return Some("当前没有注册任何工具".into());
|
||||||
}
|
}
|
||||||
|
|||||||
+66
-31
@@ -30,18 +30,21 @@ use std::sync::Arc;
|
|||||||
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
|
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
|
||||||
use tokio::sync::Mutex;
|
use tokio::sync::Mutex;
|
||||||
|
|
||||||
|
/// ## 工具执行器返回值
|
||||||
|
///
|
||||||
|
/// `Ready` 表示工具结果已经可直接写入对话;`Pending` 表示工具调用已交给
|
||||||
|
/// 外部队列,结果稍后通过 `ToolResult` 回喂。
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||||
|
pub enum ToolExecution {
|
||||||
|
Ready(String),
|
||||||
|
Pending,
|
||||||
|
}
|
||||||
|
|
||||||
/// ## 工具执行器 —— LLM 调用的工具回调
|
/// ## 工具执行器 —— LLM 调用的工具回调
|
||||||
///
|
///
|
||||||
/// 这是一个函数别名:`Fn(工具名, JSON参数) → Future<Result<String>>`。
|
|
||||||
///
|
|
||||||
/// 当 LLM 在对话中发起工具调用时,Conversation 会调用这个 executor 来执行工具。
|
/// 当 LLM 在对话中发起工具调用时,Conversation 会调用这个 executor 来执行工具。
|
||||||
///
|
|
||||||
/// ### 两种执行模式
|
|
||||||
/// 1. **同步模式** — executor 直接执行并返回结果字符串
|
|
||||||
/// 2. **异步模式** — executor 将工具调用入队到消息队列,返回 `ASYNC_MARKER + 工具名`
|
|
||||||
/// 表示"工具已异步入队,结果稍后通过 ToolResult 消息回喂"
|
|
||||||
pub type ToolExecutor = Arc<
|
pub type ToolExecutor = Arc<
|
||||||
dyn Fn(&str, &str, &str) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send>>
|
dyn Fn(&str, &str, &str) -> Pin<Box<dyn Future<Output = Result<ToolExecution, String>> + Send>>
|
||||||
+ Send
|
+ Send
|
||||||
+ Sync,
|
+ Sync,
|
||||||
>;
|
>;
|
||||||
@@ -74,9 +77,6 @@ pub struct ChatResult {
|
|||||||
pub has_pending_async: bool,
|
pub has_pending_async: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// 异步工具标记 —— 工具执行器返回此前缀表示工具已异步入队
|
|
||||||
pub const ASYNC_MARKER: &str = "__ASYNC_PENDING__";
|
|
||||||
|
|
||||||
/// ## Conversation —— LLM 对话管理器(核心类)
|
/// ## Conversation —— LLM 对话管理器(核心类)
|
||||||
///
|
///
|
||||||
/// 这是整个项目中"与 LLM 对话"的核心抽象,负责:
|
/// 这是整个项目中"与 LLM 对话"的核心抽象,负责:
|
||||||
@@ -262,19 +262,7 @@ impl Conversation {
|
|||||||
|
|
||||||
/// 通知一个异步工具结果已返回,返回 true 表示所有异步工具均已返回
|
/// 通知一个异步工具结果已返回,返回 true 表示所有异步工具均已返回
|
||||||
pub fn notify_tool_result(&self) -> bool {
|
pub fn notify_tool_result(&self) -> bool {
|
||||||
loop {
|
decrement_pending_count(&self.pending_async_count)
|
||||||
let current = self.pending_async_count.load(Ordering::SeqCst);
|
|
||||||
if current == 0 {
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
if self
|
|
||||||
.pending_async_count
|
|
||||||
.compare_exchange(current, current - 1, Ordering::SeqCst, Ordering::SeqCst)
|
|
||||||
.is_ok()
|
|
||||||
{
|
|
||||||
return current == 1;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// 工具循环核心:调 LLM → 处理工具调用 → 循环直到无工具或需要异步等待
|
/// 工具循环核心:调 LLM → 处理工具调用 → 循环直到无工具或需要异步等待
|
||||||
@@ -347,17 +335,20 @@ impl Conversation {
|
|||||||
Some(exec) => {
|
Some(exec) => {
|
||||||
match exec(&tc.function.name, &tc.function.arguments, &tc.id).await {
|
match exec(&tc.function.name, &tc.function.arguments, &tc.id).await {
|
||||||
Ok(r) => r,
|
Ok(r) => r,
|
||||||
Err(e) => format!("执行失败: {}", e),
|
Err(e) => ToolExecution::Ready(format!("执行失败: {}", e)),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
None => "工具执行器未配置".to_string(),
|
None => ToolExecution::Ready("工具执行器未配置".to_string()),
|
||||||
};
|
};
|
||||||
|
|
||||||
if result.starts_with(ASYNC_MARKER) {
|
let result = match result {
|
||||||
has_async = true;
|
ToolExecution::Ready(result) => result,
|
||||||
async_count += 1;
|
ToolExecution::Pending => {
|
||||||
continue;
|
has_async = true;
|
||||||
}
|
async_count += 1;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
tracing::info!(target: "ias::msg", "[tool] {}: {}", tc.function.name, result);
|
tracing::info!(target: "ias::msg", "[tool] {}: {}", tc.function.name, result);
|
||||||
self.session
|
self.session
|
||||||
@@ -408,6 +399,21 @@ impl Conversation {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn decrement_pending_count(count: &AtomicU32) -> bool {
|
||||||
|
loop {
|
||||||
|
let current = count.load(Ordering::SeqCst);
|
||||||
|
if current == 0 {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
if count
|
||||||
|
.compare_exchange(current, current - 1, Ordering::SeqCst, Ordering::SeqCst)
|
||||||
|
.is_ok()
|
||||||
|
{
|
||||||
|
return current == 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// ─── ChatHandle ───
|
// ─── ChatHandle ───
|
||||||
|
|
||||||
#[allow(dead_code)]
|
#[allow(dead_code)]
|
||||||
@@ -474,3 +480,32 @@ pub const DEFAULT_SYSTEM_PROMPT: &str = "\
|
|||||||
\n\
|
\n\
|
||||||
流程:收到消息 → query_capabilities → call_capability → 回复\n\
|
流程:收到消息 → query_capabilities → call_capability → 回复\n\
|
||||||
";
|
";
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn tool_execution_is_explicit_state() {
|
||||||
|
assert_eq!(
|
||||||
|
ToolExecution::Ready("__ASYNC_PENDING__weather".to_string()),
|
||||||
|
ToolExecution::Ready("__ASYNC_PENDING__weather".to_string())
|
||||||
|
);
|
||||||
|
assert_ne!(
|
||||||
|
ToolExecution::Ready("__ASYNC_PENDING__weather".to_string()),
|
||||||
|
ToolExecution::Pending
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn notify_tool_result_does_not_underflow() {
|
||||||
|
let count = AtomicU32::new(0);
|
||||||
|
|
||||||
|
assert!(!decrement_pending_count(&count));
|
||||||
|
|
||||||
|
count.store(1, Ordering::SeqCst);
|
||||||
|
assert!(decrement_pending_count(&count));
|
||||||
|
assert!(!decrement_pending_count(&count));
|
||||||
|
assert_eq!(count.load(Ordering::SeqCst), 0);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
+1
-1
@@ -22,6 +22,6 @@ pub mod provider;
|
|||||||
pub mod types;
|
pub mod types;
|
||||||
|
|
||||||
pub use conversation::{
|
pub use conversation::{
|
||||||
ASYNC_MARKER, ChatResult, Conversation, DEFAULT_SYSTEM_PROMPT, ToolExecutor,
|
ChatResult, Conversation, DEFAULT_SYSTEM_PROMPT, ToolExecution, ToolExecutor,
|
||||||
};
|
};
|
||||||
pub use types::ConversationConfig;
|
pub use types::ConversationConfig;
|
||||||
|
|||||||
@@ -12,6 +12,7 @@ use crate::tools::executor;
|
|||||||
use crate::tools::parser;
|
use crate::tools::parser;
|
||||||
use crate::tools::spec::ToolSpec;
|
use crate::tools::spec::ToolSpec;
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
|
use tokio::sync::RwLock;
|
||||||
|
|
||||||
#[derive(Debug, Clone, serde::Serialize)]
|
#[derive(Debug, Clone, serde::Serialize)]
|
||||||
pub struct ToolSummary {
|
pub struct ToolSummary {
|
||||||
@@ -53,6 +54,11 @@ pub struct Registry {
|
|||||||
tools: HashMap<String, ToolSpec>,
|
tools: HashMap<String, ToolSpec>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub struct RegistryManager {
|
||||||
|
root: String,
|
||||||
|
registry: RwLock<Registry>,
|
||||||
|
}
|
||||||
|
|
||||||
impl Registry {
|
impl Registry {
|
||||||
/// 扫描工具根目录下各子目录的 `specs/*.tool.yaml`。
|
/// 扫描工具根目录下各子目录的 `specs/*.tool.yaml`。
|
||||||
pub fn load(dir: &str) -> Self {
|
pub fn load(dir: &str) -> Self {
|
||||||
@@ -176,6 +182,68 @@ impl Registry {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
impl RegistryManager {
|
||||||
|
pub fn load(dir: impl Into<String>) -> Self {
|
||||||
|
let root = dir.into();
|
||||||
|
Self {
|
||||||
|
registry: RwLock::new(Registry::load(&root)),
|
||||||
|
root,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn reload(&self) -> usize {
|
||||||
|
let registry = Registry::load(&self.root);
|
||||||
|
let len = registry.len();
|
||||||
|
*self.registry.write().await = registry;
|
||||||
|
len
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn list_tools(&self) -> Vec<ToolSummary> {
|
||||||
|
self.registry.read().await.list_tools()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn dispatch(
|
||||||
|
&self,
|
||||||
|
name: &str,
|
||||||
|
params: &Value,
|
||||||
|
user_id: &str,
|
||||||
|
approved: bool,
|
||||||
|
approval: &Arc<ApprovalManager>,
|
||||||
|
memory: &Arc<MemoryStore>,
|
||||||
|
wechat_send: &(
|
||||||
|
dyn Fn(&str, &str) -> std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send>>
|
||||||
|
+ Send
|
||||||
|
+ Sync
|
||||||
|
),
|
||||||
|
) -> String {
|
||||||
|
let spec = {
|
||||||
|
let registry = self.registry.read().await;
|
||||||
|
match registry.get(name) {
|
||||||
|
Some(spec) => spec.clone(),
|
||||||
|
None => {
|
||||||
|
tracing::warn!(target:"ias::registry","未知: {name}");
|
||||||
|
return format!("未知工具: {name}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
if let Some(err) = validate_required_params(&spec, params) {
|
||||||
|
return err;
|
||||||
|
}
|
||||||
|
|
||||||
|
executor::execute_loop(
|
||||||
|
&spec,
|
||||||
|
params,
|
||||||
|
user_id,
|
||||||
|
approved,
|
||||||
|
approval,
|
||||||
|
memory,
|
||||||
|
wechat_send,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn validate_required_params(spec: &ToolSpec, params: &Value) -> Option<String> {
|
fn validate_required_params(spec: &ToolSpec, params: &Value) -> Option<String> {
|
||||||
let Some(obj) = params.as_object() else {
|
let Some(obj) = params.as_object() else {
|
||||||
return Some(format!("工具 '{}' 参数必须是 JSON 对象", spec.name));
|
return Some(format!("工具 '{}' 参数必须是 JSON 对象", spec.name));
|
||||||
|
|||||||
+3
-3
@@ -2,10 +2,10 @@
|
|||||||
//!
|
//!
|
||||||
//! 与 `*.tool.yaml` 格式一一对应的 Rust 结构体。
|
//! 与 `*.tool.yaml` 格式一一对应的 Rust 结构体。
|
||||||
|
|
||||||
use serde::Deserialize;
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
/// 单个参数的声明
|
/// 单个参数的声明
|
||||||
#[derive(Debug, Clone, Deserialize)]
|
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||||
pub struct ParamSpec {
|
pub struct ParamSpec {
|
||||||
pub name: String,
|
pub name: String,
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
@@ -30,7 +30,7 @@ pub struct ParamSpec {
|
|||||||
/// required: true
|
/// required: true
|
||||||
/// desc: 城市名称
|
/// desc: 城市名称
|
||||||
/// ```
|
/// ```
|
||||||
#[derive(Debug, Clone, Deserialize)]
|
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||||
pub struct ToolSpec {
|
pub struct ToolSpec {
|
||||||
pub name: String,
|
pub name: String,
|
||||||
pub desc: String,
|
pub desc: String,
|
||||||
|
|||||||
+1
-1
@@ -3,7 +3,7 @@
|
|||||||
set -e
|
set -e
|
||||||
cd "$(dirname "$0")"
|
cd "$(dirname "$0")"
|
||||||
|
|
||||||
TOOLS=(datetime weather amap web_search fetch_page memories)
|
TOOLS=(datetime weather amap web_search fetch_page memories tool_manager)
|
||||||
|
|
||||||
echo "=== 构建工具 ==="
|
echo "=== 构建工具 ==="
|
||||||
for tool in "${TOOLS[@]}"; do
|
for tool in "${TOOLS[@]}"; do
|
||||||
|
|||||||
Generated
+105
@@ -0,0 +1,105 @@
|
|||||||
|
# This file is automatically @generated by Cargo.
|
||||||
|
# It is not intended for manual editing.
|
||||||
|
version = 4
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "itoa"
|
||||||
|
version = "1.0.18"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "memchr"
|
||||||
|
version = "2.8.2"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "proc-macro2"
|
||||||
|
version = "1.0.106"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934"
|
||||||
|
dependencies = [
|
||||||
|
"unicode-ident",
|
||||||
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "quote"
|
||||||
|
version = "1.0.45"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924"
|
||||||
|
dependencies = [
|
||||||
|
"proc-macro2",
|
||||||
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "serde"
|
||||||
|
version = "1.0.228"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e"
|
||||||
|
dependencies = [
|
||||||
|
"serde_core",
|
||||||
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "serde_core"
|
||||||
|
version = "1.0.228"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad"
|
||||||
|
dependencies = [
|
||||||
|
"serde_derive",
|
||||||
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "serde_derive"
|
||||||
|
version = "1.0.228"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79"
|
||||||
|
dependencies = [
|
||||||
|
"proc-macro2",
|
||||||
|
"quote",
|
||||||
|
"syn",
|
||||||
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "serde_json"
|
||||||
|
version = "1.0.150"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9"
|
||||||
|
dependencies = [
|
||||||
|
"itoa",
|
||||||
|
"memchr",
|
||||||
|
"serde",
|
||||||
|
"serde_core",
|
||||||
|
"zmij",
|
||||||
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "syn"
|
||||||
|
version = "2.0.118"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422"
|
||||||
|
dependencies = [
|
||||||
|
"proc-macro2",
|
||||||
|
"quote",
|
||||||
|
"unicode-ident",
|
||||||
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "tool_manager"
|
||||||
|
version = "0.1.0"
|
||||||
|
dependencies = [
|
||||||
|
"serde_json",
|
||||||
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "unicode-ident"
|
||||||
|
version = "1.0.24"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "zmij"
|
||||||
|
version = "1.0.21"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa"
|
||||||
@@ -0,0 +1,11 @@
|
|||||||
|
[package]
|
||||||
|
name = "tool_manager"
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2024"
|
||||||
|
|
||||||
|
[[bin]]
|
||||||
|
name = "tool_manager"
|
||||||
|
path = "src/main.rs"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
serde_json = "1.0"
|
||||||
@@ -0,0 +1,32 @@
|
|||||||
|
name: tool_manager
|
||||||
|
desc: 创建、修改和构建本地能力工具。用于给助手新增或更新 tools 目录下的 stdio Rust 工具;高风险操作,执行前需要用户确认。
|
||||||
|
type: stdio
|
||||||
|
tool: tool_manager
|
||||||
|
path: /target/release/tool_manager
|
||||||
|
risk_level: high
|
||||||
|
timeout_secs: 120
|
||||||
|
params:
|
||||||
|
- name: action
|
||||||
|
required: true
|
||||||
|
desc: 操作类型:create_tool、update_tool 或 build_tool
|
||||||
|
- name: tool_name
|
||||||
|
required: true
|
||||||
|
desc: 工具名称,只允许字母、数字、下划线和短横线
|
||||||
|
- name: description
|
||||||
|
required: false
|
||||||
|
desc: 工具说明,创建或更新 spec 时使用
|
||||||
|
- name: rust_code
|
||||||
|
required: false
|
||||||
|
desc: 完整的 src/main.rs Rust 源码;省略时 create_tool 会生成默认模板
|
||||||
|
- name: params
|
||||||
|
required: false
|
||||||
|
desc: 参数列表数组,每项包含 name、required、desc
|
||||||
|
- name: risk_level
|
||||||
|
required: false
|
||||||
|
desc: 新工具风险级别,low 或 high,默认 low
|
||||||
|
- name: timeout_secs
|
||||||
|
required: false
|
||||||
|
desc: 新工具超时时间秒数,默认 30
|
||||||
|
- name: overwrite
|
||||||
|
required: false
|
||||||
|
desc: create_tool 目标已存在时是否覆盖,默认 false
|
||||||
@@ -0,0 +1,258 @@
|
|||||||
|
//! tool_manager — 受控创建、修改、构建本地工具
|
||||||
|
|
||||||
|
use serde_json::{Value, json};
|
||||||
|
use std::fs;
|
||||||
|
use std::path::{Path, PathBuf};
|
||||||
|
use std::process::Command;
|
||||||
|
|
||||||
|
fn main() {
|
||||||
|
let input: Value = match serde_json::from_reader(std::io::stdin()) {
|
||||||
|
Ok(v) => v,
|
||||||
|
Err(e) => return result(&format!("读取输入失败: {e}")),
|
||||||
|
};
|
||||||
|
|
||||||
|
let params = &input["params"];
|
||||||
|
let action = params["action"].as_str().unwrap_or("");
|
||||||
|
let tool_name = params["tool_name"].as_str().unwrap_or("");
|
||||||
|
|
||||||
|
let output = match run(action, tool_name, params) {
|
||||||
|
Ok(msg) => msg,
|
||||||
|
Err(e) => format!("工具管理失败: {e}"),
|
||||||
|
};
|
||||||
|
result(&output);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn run(action: &str, tool_name: &str, params: &Value) -> Result<String, String> {
|
||||||
|
validate_tool_name(tool_name)?;
|
||||||
|
let tools_root = locate_tools_root()?;
|
||||||
|
let tool_dir = tools_root.join(tool_name);
|
||||||
|
|
||||||
|
match action {
|
||||||
|
"create_tool" => {
|
||||||
|
let overwrite = truthy(params.get("overwrite"));
|
||||||
|
if tool_dir.exists() && !overwrite {
|
||||||
|
return Err(format!(
|
||||||
|
"工具 {tool_name} 已存在;如需覆盖,请设置 overwrite=true"
|
||||||
|
));
|
||||||
|
}
|
||||||
|
write_tool_project(&tool_dir, tool_name, params)?;
|
||||||
|
build_tool(&tool_dir, tool_name)?;
|
||||||
|
Ok(format!(
|
||||||
|
"已创建并构建工具 {tool_name}。可通过 query_capabilities 查看并用 call_capability 调用。"
|
||||||
|
))
|
||||||
|
}
|
||||||
|
"update_tool" => {
|
||||||
|
if !tool_dir.exists() {
|
||||||
|
return Err(format!("工具 {tool_name} 不存在,无法修改"));
|
||||||
|
}
|
||||||
|
update_tool_project(&tool_dir, tool_name, params)?;
|
||||||
|
build_tool(&tool_dir, tool_name)?;
|
||||||
|
Ok(format!(
|
||||||
|
"已修改并构建工具 {tool_name}。注册表重载后即可使用最新版本。"
|
||||||
|
))
|
||||||
|
}
|
||||||
|
"build_tool" => {
|
||||||
|
if !tool_dir.exists() {
|
||||||
|
return Err(format!("工具 {tool_name} 不存在,无法构建"));
|
||||||
|
}
|
||||||
|
build_tool(&tool_dir, tool_name)?;
|
||||||
|
Ok(format!("已构建工具 {tool_name}。"))
|
||||||
|
}
|
||||||
|
_ => Err("action 必须是 create_tool、update_tool 或 build_tool".to_string()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn write_tool_project(tool_dir: &Path, tool_name: &str, params: &Value) -> Result<(), String> {
|
||||||
|
fs::create_dir_all(tool_dir.join("src")).map_err(|e| e.to_string())?;
|
||||||
|
fs::create_dir_all(tool_dir.join("specs")).map_err(|e| e.to_string())?;
|
||||||
|
fs::write(tool_dir.join("Cargo.toml"), cargo_toml(tool_name)).map_err(|e| e.to_string())?;
|
||||||
|
fs::write(tool_dir.join("src/main.rs"), rust_code(tool_name, params))
|
||||||
|
.map_err(|e| e.to_string())?;
|
||||||
|
fs::write(
|
||||||
|
tool_dir
|
||||||
|
.join("specs")
|
||||||
|
.join(format!("{tool_name}.tool.yaml")),
|
||||||
|
spec_yaml(tool_name, params),
|
||||||
|
)
|
||||||
|
.map_err(|e| e.to_string())?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn update_tool_project(tool_dir: &Path, tool_name: &str, params: &Value) -> Result<(), String> {
|
||||||
|
if let Some(code) = params["rust_code"].as_str() {
|
||||||
|
fs::write(tool_dir.join("src/main.rs"), code).map_err(|e| e.to_string())?;
|
||||||
|
}
|
||||||
|
|
||||||
|
let should_update_spec = params.get("description").is_some()
|
||||||
|
|| params.get("params").is_some()
|
||||||
|
|| params.get("risk_level").is_some()
|
||||||
|
|| params.get("timeout_secs").is_some();
|
||||||
|
if should_update_spec {
|
||||||
|
fs::create_dir_all(tool_dir.join("specs")).map_err(|e| e.to_string())?;
|
||||||
|
fs::write(
|
||||||
|
tool_dir
|
||||||
|
.join("specs")
|
||||||
|
.join(format!("{tool_name}.tool.yaml")),
|
||||||
|
spec_yaml(tool_name, params),
|
||||||
|
)
|
||||||
|
.map_err(|e| e.to_string())?;
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn build_tool(tool_dir: &Path, tool_name: &str) -> Result<(), String> {
|
||||||
|
let output = Command::new("cargo")
|
||||||
|
.arg("build")
|
||||||
|
.arg("--release")
|
||||||
|
.current_dir(tool_dir)
|
||||||
|
.output()
|
||||||
|
.map_err(|e| format!("启动 cargo 失败: {e}"))?;
|
||||||
|
|
||||||
|
if !output.status.success() {
|
||||||
|
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||||
|
return Err(format!("cargo build 失败: {}", stderr.trim()));
|
||||||
|
}
|
||||||
|
|
||||||
|
let binary = tool_dir.join("target/release").join(tool_name);
|
||||||
|
if !binary.exists() {
|
||||||
|
return Err(format!("构建后未找到二进制: {}", binary.display()));
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn locate_tools_root() -> Result<PathBuf, String> {
|
||||||
|
let exe = std::env::current_exe().map_err(|e| e.to_string())?;
|
||||||
|
let tool_dir = exe
|
||||||
|
.parent()
|
||||||
|
.and_then(Path::parent)
|
||||||
|
.and_then(Path::parent)
|
||||||
|
.ok_or_else(|| "无法定位 tool_manager 目录".to_string())?;
|
||||||
|
tool_dir
|
||||||
|
.parent()
|
||||||
|
.map(Path::to_path_buf)
|
||||||
|
.ok_or_else(|| "无法定位 tools 根目录".to_string())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn cargo_toml(tool_name: &str) -> String {
|
||||||
|
format!(
|
||||||
|
r#"[package]
|
||||||
|
name = "{tool_name}"
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2024"
|
||||||
|
|
||||||
|
[[bin]]
|
||||||
|
name = "{tool_name}"
|
||||||
|
path = "src/main.rs"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
serde_json = "1.0"
|
||||||
|
"#
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn rust_code(tool_name: &str, params: &Value) -> String {
|
||||||
|
params["rust_code"]
|
||||||
|
.as_str()
|
||||||
|
.map(str::to_string)
|
||||||
|
.unwrap_or_else(|| default_rust_code(tool_name))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_rust_code(tool_name: &str) -> String {
|
||||||
|
format!(
|
||||||
|
r#"//! {tool_name} — generated tool
|
||||||
|
|
||||||
|
fn main() {{
|
||||||
|
let input: serde_json::Value = serde_json::from_reader(std::io::stdin()).unwrap_or_default();
|
||||||
|
let content = format!("工具 {tool_name} 已收到参数: {{}}", input["params"]);
|
||||||
|
println!("{{}}", serde_json::json!({{"type":"result","content":content}}));
|
||||||
|
}}
|
||||||
|
"#
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn spec_yaml(tool_name: &str, params: &Value) -> String {
|
||||||
|
let desc = params["description"]
|
||||||
|
.as_str()
|
||||||
|
.unwrap_or("由 tool_manager 创建的本地工具");
|
||||||
|
let risk_level = match params["risk_level"].as_str() {
|
||||||
|
Some("high") => "high",
|
||||||
|
_ => "low",
|
||||||
|
};
|
||||||
|
let timeout_secs = params["timeout_secs"].as_u64().unwrap_or(30).clamp(1, 300);
|
||||||
|
|
||||||
|
let mut out = format!(
|
||||||
|
"name: {tool_name}\n\
|
||||||
|
desc: {}\n\
|
||||||
|
type: stdio\n\
|
||||||
|
tool: {tool_name}\n\
|
||||||
|
path: /target/release/{tool_name}\n\
|
||||||
|
risk_level: {risk_level}\n\
|
||||||
|
timeout_secs: {timeout_secs}\n",
|
||||||
|
yaml_scalar(desc)
|
||||||
|
);
|
||||||
|
|
||||||
|
if let Some(items) = params["params"].as_array() {
|
||||||
|
out.push_str("params:\n");
|
||||||
|
for item in items {
|
||||||
|
let Some(name) = item["name"].as_str() else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
if validate_param_name(name).is_err() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let required = item["required"].as_bool().unwrap_or(false);
|
||||||
|
let desc = item["desc"].as_str().unwrap_or("");
|
||||||
|
out.push_str(&format!(
|
||||||
|
" - name: {name}\n required: {required}\n desc: {}\n",
|
||||||
|
yaml_scalar(desc)
|
||||||
|
));
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
out.push_str("params: []\n");
|
||||||
|
}
|
||||||
|
|
||||||
|
out
|
||||||
|
}
|
||||||
|
|
||||||
|
fn validate_tool_name(name: &str) -> Result<(), String> {
|
||||||
|
if name.is_empty() {
|
||||||
|
return Err("缺少 tool_name".to_string());
|
||||||
|
}
|
||||||
|
if matches!(name, "target" | "specs" | "src" | "." | "..") {
|
||||||
|
return Err("tool_name 使用了保留名称".to_string());
|
||||||
|
}
|
||||||
|
if !name
|
||||||
|
.chars()
|
||||||
|
.all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
|
||||||
|
{
|
||||||
|
return Err("tool_name 只允许字母、数字、下划线和短横线".to_string());
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn validate_param_name(name: &str) -> Result<(), String> {
|
||||||
|
if name.is_empty()
|
||||||
|
|| !name
|
||||||
|
.chars()
|
||||||
|
.all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
|
||||||
|
{
|
||||||
|
return Err("参数名非法".to_string());
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn yaml_scalar(value: &str) -> String {
|
||||||
|
serde_json::to_string(value).unwrap_or_else(|_| "\"\"".to_string())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn truthy(value: Option<&Value>) -> bool {
|
||||||
|
match value {
|
||||||
|
Some(Value::Bool(v)) => *v,
|
||||||
|
Some(Value::String(v)) => matches!(v.as_str(), "true" | "1" | "yes" | "是"),
|
||||||
|
_ => false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn result(content: &str) {
|
||||||
|
println!("{}", json!({"type":"result","content":content}));
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user