Refactor tool execution pipeline
This commit is contained in:
@@ -0,0 +1,304 @@
|
||||
//! # 工具注册器
|
||||
//!
|
||||
//! 扫描每个工具目录下的 `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;
|
||||
|
||||
#[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>,
|
||||
}
|
||||
|
||||
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()
|
||||
}
|
||||
}
|
||||
|
||||
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()
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user