373 lines
10 KiB
Rust
373 lines
10 KiB
Rust
//! # 工具注册器
|
||
//!
|
||
//! 扫描每个工具目录下的 `specs/*.tool.yaml`，按 `name` 路由分发。
|
||
|
||
use std::collections::HashMap;
|
||
use std::path::{Path, PathBuf};
|
||
use std::sync::Arc;
|
||
|
||
use crate::context::MemoryStore;
|
||
use crate::tools::approval::ApprovalManager;
|
||
use crate::tools::executor;
|
||
use crate::tools::parser;
|
||
use crate::tools::spec::ToolSpec;
|
||
use serde_json::Value;
|
||
use tokio::sync::RwLock;
|
||
|
||
#[derive(Debug, Clone, serde::Serialize)]
|
||
pub struct ToolSummary {
|
||
pub name: String,
|
||
pub desc: String,
|
||
pub params: Vec<ParamSummary>,
|
||
}
|
||
|
||
#[derive(Debug, Clone, serde::Serialize)]
|
||
pub struct ParamSummary {
|
||
pub name: String,
|
||
pub required: bool,
|
||
pub desc: String,
|
||
}
|
||
|
||
/// 解析工具目录。相对路径基于二进制所在目录。
|
||
fn resolve_dir(dir: &str) -> PathBuf {
|
||
let path = Path::new(dir);
|
||
if path.is_absolute() {
|
||
return path.to_path_buf();
|
||
}
|
||
let exe_dir = std::env::current_exe()
|
||
.ok()
|
||
.and_then(|e| e.parent().map(|p| p.to_path_buf()))
|
||
.unwrap_or_else(|| PathBuf::from("."));
|
||
let candidate = exe_dir.join(dir);
|
||
if candidate.exists() {
|
||
return candidate;
|
||
}
|
||
let dev = exe_dir.join("../..").join(dir);
|
||
if dev.exists() {
|
||
tracing::info!(target:"ias::registry","dev模式: {}", dev.display());
|
||
return dev;
|
||
}
|
||
path.to_path_buf()
|
||
}
|
||
|
||
pub struct Registry {
|
||
tools: HashMap<String, ToolSpec>,
|
||
}
|
||
|
||
pub struct RegistryManager {
|
||
root: String,
|
||
registry: RwLock<Registry>,
|
||
}
|
||
|
||
impl Registry {
|
||
/// 扫描工具根目录下各子目录的 `specs/*.tool.yaml`。
|
||
pub fn load(dir: &str) -> Self {
|
||
let root = resolve_dir(dir);
|
||
let mut tools = HashMap::new();
|
||
|
||
let entries = match std::fs::read_dir(&root) {
|
||
Ok(e) => e,
|
||
Err(e) => {
|
||
tracing::warn!(target:"ias::registry","读取目录失败 {}: {}", root.display(), e);
|
||
return Self { tools };
|
||
}
|
||
};
|
||
|
||
for entry in entries.flatten() {
|
||
let path = entry.path();
|
||
if !path.is_dir() {
|
||
continue;
|
||
}
|
||
let specs_dir = path.join("specs");
|
||
if !specs_dir.exists() {
|
||
continue;
|
||
}
|
||
|
||
let result = parser::parse_dir(&specs_dir);
|
||
for (file, err) in &result.errors {
|
||
tracing::warn!(target:"ias::registry","解析失败 {}: {}", file, err);
|
||
}
|
||
|
||
for mut spec in result.specs {
|
||
// 解析二进制路径
|
||
let binary = match &spec.path {
|
||
Some(p) => {
|
||
let p = p.trim_start_matches('/');
|
||
path.join(p)
|
||
}
|
||
None => {
|
||
// 默认: {tool_dir}/target/release/{tool}
|
||
path.join("target/release").join(&spec.tool)
|
||
}
|
||
};
|
||
|
||
if !binary.exists() {
|
||
tracing::warn!(target:"ias::registry","{} 二进制不存在: {}", spec.name, binary.display());
|
||
continue;
|
||
}
|
||
|
||
// 将绝对路径存入 path 字段（executor 直接使用）
|
||
spec.path = Some(binary.to_string_lossy().to_string());
|
||
|
||
tracing::info!(target:"ias::registry","注册: {} ({})", spec.name, spec.r#type);
|
||
tools.insert(spec.name.clone(), spec);
|
||
}
|
||
}
|
||
|
||
Self { tools }
|
||
}
|
||
|
||
pub fn list_tools(&self) -> Vec<ToolSummary> {
|
||
self.tools
|
||
.values()
|
||
.map(|s| ToolSummary {
|
||
name: s.name.clone(),
|
||
desc: s.desc.clone(),
|
||
params: s
|
||
.params
|
||
.iter()
|
||
.map(|p| ParamSummary {
|
||
name: p.name.clone(),
|
||
required: p.required,
|
||
desc: p.desc.clone(),
|
||
})
|
||
.collect(),
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
pub fn get(&self, name: &str) -> Option<&ToolSpec> {
|
||
self.tools.get(name)
|
||
}
|
||
|
||
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 = match self.get(name) {
|
||
Some(s) => s,
|
||
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
|
||
}
|
||
|
||
#[allow(dead_code)]
|
||
pub fn len(&self) -> usize {
|
||
self.tools.len()
|
||
}
|
||
}
|
||
|
||
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> {
|
||
let Some(obj) = params.as_object() else {
|
||
return Some(format!("工具 '{}' 参数必须是 JSON 对象", spec.name));
|
||
};
|
||
|
||
let missing: Vec<&str> = spec
|
||
.params
|
||
.iter()
|
||
.filter(|p| p.required)
|
||
.filter_map(|p| match obj.get(&p.name) {
|
||
Some(Value::Null) | None => Some(p.name.as_str()),
|
||
Some(Value::String(s)) if s.trim().is_empty() => Some(p.name.as_str()),
|
||
Some(_) => None,
|
||
})
|
||
.collect();
|
||
|
||
if missing.is_empty() {
|
||
None
|
||
} else {
|
||
Some(format!(
|
||
"缺少必填参数: {}（工具: {}）",
|
||
missing.join(", "),
|
||
spec.name
|
||
))
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use std::io::Write;
|
||
|
||
#[test]
|
||
fn test_load_specs_dir() {
|
||
let dir = std::env::temp_dir().join("ias_test_specs_dir");
|
||
let _ = std::fs::remove_dir_all(&dir);
|
||
let tool_dir = dir.join("mytool");
|
||
std::fs::create_dir_all(tool_dir.join("specs")).unwrap();
|
||
std::fs::create_dir_all(tool_dir.join("target/release")).unwrap();
|
||
|
||
std::fs::write(
|
||
tool_dir.join("specs/test.tool.yaml"),
|
||
r#"
|
||
name: test_tool
|
||
desc: 测试
|
||
type: http
|
||
tool: mytool
|
||
path: /target/release/mytool
|
||
"#,
|
||
)
|
||
.unwrap();
|
||
|
||
let mut f = std::fs::File::create(tool_dir.join("target/release/mytool")).unwrap();
|
||
f.write_all(b"fake").unwrap();
|
||
#[cfg(unix)]
|
||
{
|
||
use std::os::unix::fs::PermissionsExt;
|
||
let _ = f.set_permissions(std::fs::Permissions::from_mode(0o755));
|
||
}
|
||
|
||
let registry = Registry::load(dir.to_str().unwrap());
|
||
assert_eq!(registry.len(), 1);
|
||
|
||
let tools = registry.list_tools();
|
||
assert_eq!(tools[0].name, "test_tool");
|
||
|
||
let _ = std::fs::remove_dir_all(&dir);
|
||
}
|
||
|
||
#[test]
|
||
fn test_load_real_tools() {
|
||
let dir = std::path::Path::new("tools");
|
||
if !dir.exists() {
|
||
return;
|
||
}
|
||
let registry = Registry::load("tools");
|
||
assert!(registry.len() >= 1, "至少 1 个工具");
|
||
}
|
||
|
||
#[test]
|
||
fn test_validate_required_params_missing() {
|
||
let spec = ToolSpec {
|
||
name: "weather_now".into(),
|
||
desc: "天气".into(),
|
||
r#type: "http".into(),
|
||
tool: "weather".into(),
|
||
command: None,
|
||
path: None,
|
||
risk_level: "low".into(),
|
||
timeout_secs: 10,
|
||
env: vec![],
|
||
params: vec![crate::tools::spec::ParamSpec {
|
||
name: "location".into(),
|
||
required: true,
|
||
desc: "城市".into(),
|
||
}],
|
||
};
|
||
|
||
let err = validate_required_params(&spec, &serde_json::json!({})).unwrap();
|
||
assert!(err.contains("location"));
|
||
}
|
||
|
||
#[test]
|
||
fn test_validate_required_params_present() {
|
||
let spec = ToolSpec {
|
||
name: "weather_now".into(),
|
||
desc: "天气".into(),
|
||
r#type: "http".into(),
|
||
tool: "weather".into(),
|
||
command: None,
|
||
path: None,
|
||
risk_level: "low".into(),
|
||
timeout_secs: 10,
|
||
env: vec![],
|
||
params: vec![crate::tools::spec::ParamSpec {
|
||
name: "location".into(),
|
||
required: true,
|
||
desc: "城市".into(),
|
||
}],
|
||
};
|
||
|
||
assert!(
|
||
validate_required_params(&spec, &serde_json::json!({"location": "北京"})).is_none()
|
||
);
|
||
}
|
||
}
|