refactor: worker进程 → tokio task 架构 + 多渠道消息队列 + JSON工具系统
核心变更: - daemon: 废弃 worker 子进程模型,替换为 LLM/Tool/Send 三消费者 tokio task - queue: 新增公平轮转 MessageQueue + QueueRunner 路由,严格 ChannelId 渠道隔离 - channel: ChannelId(platform + user_id) 统一渠道标识,替代散列 user_id+platform - tools: 全部工具统一经由 execute_tool() → Tool Consumer 单一执行路径 - llm: chat_with_tools 返回 ChatResult,支持异步工具调用 → 保留/恢复 session - tools: 新增 JSON 工具定义系统 (ToolDef) 及 SubprocessRunner (HTTP 模板引擎) - bin: ias-auth 独立 JWT 生成器,ias-fetch headless Chrome 网页抓取 - 所有出站消息通过 Send Consumer 队列转发,无直接 send_text 调用 - Session TTL 清理 (5分钟),QueueRunner 通道容错 - cargo check 零警告,cargo test 23/23 通过
This commit is contained in:
@@ -0,0 +1,10 @@
|
|||||||
|
[package]
|
||||||
|
name = "ias-fetch"
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2024"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
chromiumoxide = { version = "0.7", default-features = false, features = ["tokio-runtime", "bytes"] }
|
||||||
|
futures = "0.3"
|
||||||
|
serde_json = "1.0"
|
||||||
|
tokio = { version = "1.0", features = ["full"] }
|
||||||
@@ -0,0 +1,143 @@
|
|||||||
|
//! ias-fetch — 通过 headless Chrome 提取网页正文
|
||||||
|
//!
|
||||||
|
//! 启动无头 Chrome 渲染页面 JS,取渲染后的纯文本输出到 stdout。
|
||||||
|
//! 专供 JSON 工具 `${cmd:ias-fetch {{url}}}` 调用。
|
||||||
|
//!
|
||||||
|
//! 用法:
|
||||||
|
//! ias-fetch <url> # 默认:取 body.innerText
|
||||||
|
//! ias-fetch --raw <url> # 返回完整 HTML
|
||||||
|
//! ias-fetch --timeout 15 <url> # 自定义超时(默认 15s)
|
||||||
|
|
||||||
|
use futures::StreamExt;
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
use chromiumoxide::{Browser, BrowserConfig, Page};
|
||||||
|
|
||||||
|
const DEFAULT_TIMEOUT: u64 = 15;
|
||||||
|
const MAX_OUTPUT_CHARS: usize = 10_000;
|
||||||
|
|
||||||
|
#[tokio::main]
|
||||||
|
async fn main() {
|
||||||
|
let args: Vec<String> = std::env::args().collect();
|
||||||
|
|
||||||
|
let mut raw_mode = false;
|
||||||
|
let mut timeout_secs = DEFAULT_TIMEOUT;
|
||||||
|
let mut url: Option<String> = None;
|
||||||
|
|
||||||
|
let mut i = 1;
|
||||||
|
while i < args.len() {
|
||||||
|
match args[i].as_str() {
|
||||||
|
"--raw" => raw_mode = true,
|
||||||
|
"--timeout" => {
|
||||||
|
i += 1;
|
||||||
|
timeout_secs = args.get(i).and_then(|s| s.parse().ok()).unwrap_or(DEFAULT_TIMEOUT);
|
||||||
|
}
|
||||||
|
s if s.starts_with("http://") || s.starts_with("https://") => {
|
||||||
|
url = Some(s.to_string());
|
||||||
|
}
|
||||||
|
s if !s.starts_with('-') => url = Some(s.to_string()),
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
i += 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
let url = match url {
|
||||||
|
Some(u) => {
|
||||||
|
if u.starts_with("http://") || u.starts_with("https://") {
|
||||||
|
u
|
||||||
|
} else {
|
||||||
|
format!("https://{}", u)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
None => {
|
||||||
|
eprintln!("用法: ias-fetch [--raw] [--timeout N] <url>");
|
||||||
|
std::process::exit(1);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let timeout = Duration::from_secs(timeout_secs);
|
||||||
|
|
||||||
|
match tokio::time::timeout(timeout, fetch_page(&url, raw_mode)).await {
|
||||||
|
Ok(Ok(output)) => {
|
||||||
|
if output.is_empty() {
|
||||||
|
eprintln!("页面无内容");
|
||||||
|
std::process::exit(1);
|
||||||
|
}
|
||||||
|
println!("{}", output);
|
||||||
|
}
|
||||||
|
Ok(Err(e)) => {
|
||||||
|
eprintln!("{}", e);
|
||||||
|
std::process::exit(1);
|
||||||
|
}
|
||||||
|
Err(_) => {
|
||||||
|
eprintln!("超时 ({}s): {}", timeout_secs, url);
|
||||||
|
std::process::exit(1);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn fetch_page(url: &str, raw_mode: bool) -> Result<String, String> {
|
||||||
|
let (mut browser, mut handler) = Browser::launch(
|
||||||
|
BrowserConfig::builder()
|
||||||
|
.build()
|
||||||
|
.map_err(|e| format!("Chrome 配置失败: {}", e))?,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("启动 Chrome 失败: {}\n请确保已安装 chromium 或 google-chrome", e))?;
|
||||||
|
|
||||||
|
let handle = tokio::spawn(async move {
|
||||||
|
while let Some(_event) = handler.next().await {
|
||||||
|
// 持续消费 handler 事件
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
// new_page 接受 URL,自动导航并等待加载
|
||||||
|
let page = browser
|
||||||
|
.new_page(url)
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("导航到 {} 失败: {}", url, e))?;
|
||||||
|
|
||||||
|
// 额外等待 JS 渲染
|
||||||
|
tokio::time::sleep(Duration::from_secs(2)).await;
|
||||||
|
|
||||||
|
// 提取内容
|
||||||
|
let content = if raw_mode {
|
||||||
|
page.content().await.map_err(|e| format!("获取 HTML 失败: {}", e))?
|
||||||
|
} else {
|
||||||
|
get_inner_text(&page).await?
|
||||||
|
};
|
||||||
|
|
||||||
|
// 清理
|
||||||
|
browser.close().await.ok();
|
||||||
|
handle.await.ok();
|
||||||
|
|
||||||
|
// 截断
|
||||||
|
if content.chars().count() > MAX_OUTPUT_CHARS {
|
||||||
|
Ok(format!(
|
||||||
|
"{}…(已截断,完整内容 {} 字符)",
|
||||||
|
content.chars().take(MAX_OUTPUT_CHARS).collect::<String>(),
|
||||||
|
content.chars().count()
|
||||||
|
))
|
||||||
|
} else {
|
||||||
|
Ok(content)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 通过 evaluate JS 获取 body.innerText
|
||||||
|
async fn get_inner_text(page: &Page) -> Result<String, String> {
|
||||||
|
let result = page
|
||||||
|
.evaluate("document.body?.innerText || document.documentElement?.innerText || ''")
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("提取正文失败: {}", e))?;
|
||||||
|
|
||||||
|
let value: serde_json::Value = result
|
||||||
|
.into_value()
|
||||||
|
.map_err(|e| format!("解析 evaluate 结果失败: {}", e))?;
|
||||||
|
match value {
|
||||||
|
serde_json::Value::String(s) => {
|
||||||
|
let cleaned: Vec<&str> = s.lines().map(|l| l.trim()).filter(|l| !l.is_empty()).collect();
|
||||||
|
Ok(cleaned.join("\n"))
|
||||||
|
}
|
||||||
|
_ => Ok(String::new()),
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,104 @@
|
|||||||
|
//! ias-auth — 独立 Token 生成器
|
||||||
|
//!
|
||||||
|
//! 专供 JSON 工具定义通过 `${cmd:ias-auth gen-qweather-jwt}` 调用。
|
||||||
|
//! 编译快、无运行时依赖,独立于主程序 `ias`。
|
||||||
|
//!
|
||||||
|
//! 用法:
|
||||||
|
//! ias-auth gen-qweather-jwt 生成和风天气 Ed25519 JWT Token
|
||||||
|
|
||||||
|
use std::process::exit;
|
||||||
|
|
||||||
|
fn main() {
|
||||||
|
let args: Vec<String> = std::env::args().collect();
|
||||||
|
if args.len() < 2 {
|
||||||
|
eprintln!("用法: ias-auth <command>");
|
||||||
|
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<String, String> {
|
||||||
|
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::<Vec<_>>()
|
||||||
|
.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())
|
||||||
|
))
|
||||||
|
}
|
||||||
@@ -0,0 +1,38 @@
|
|||||||
|
//! 渠道 / 通道基础类型
|
||||||
|
//!
|
||||||
|
//! 定义消息的来源/目标渠道标识,以及从外部(微信等)收到的原始消息。
|
||||||
|
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::fmt;
|
||||||
|
|
||||||
|
/// 渠道标识 —— 区分消息来自哪个平台和哪个用户
|
||||||
|
#[derive(Debug, Clone, Hash, Eq, PartialEq, Serialize, Deserialize)]
|
||||||
|
pub struct ChannelId {
|
||||||
|
/// 平台名称,如 "wechat"、"telegram"、"console"
|
||||||
|
pub platform: String,
|
||||||
|
/// 在该平台上的用户 ID
|
||||||
|
pub user_id: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ChannelId {
|
||||||
|
/// 快捷构造微信渠道
|
||||||
|
pub fn wechat(user_id: impl Into<String>) -> Self {
|
||||||
|
Self { platform: "wechat".into(), user_id: user_id.into() }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl fmt::Display for ChannelId {
|
||||||
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||||
|
write!(f, "{}:{}", self.platform, self.user_id)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 从外部轮询收到的原始消息
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
#[allow(dead_code)]
|
||||||
|
pub struct IncomingMessage {
|
||||||
|
pub channel: ChannelId,
|
||||||
|
pub text: String,
|
||||||
|
pub message_id: String,
|
||||||
|
pub context_token: Option<String>,
|
||||||
|
}
|
||||||
+1
-1
@@ -115,7 +115,7 @@ pub enum Commands {
|
|||||||
sock: String,
|
sock: String,
|
||||||
},
|
},
|
||||||
|
|
||||||
/// Worker 进程(由 daemon 自动 spawn,也可手动测试)
|
/// [已废弃] Worker 进程(新架构使用 daemon 统一 tokio task)
|
||||||
///
|
///
|
||||||
/// 示例:
|
/// 示例:
|
||||||
/// ias worker --sock /tmp/ias_daemon.sock
|
/// ias worker --sock /tmp/ias_daemon.sock
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ use crate::llm::types::Message;
|
|||||||
use chrono::{DateTime, Utc};
|
use chrono::{DateTime, Utc};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
use uuid::Uuid;
|
||||||
|
|
||||||
/// 摘要原因
|
/// 摘要原因
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
@@ -44,6 +45,8 @@ pub struct ChatSession {
|
|||||||
pub token_budget: u32,
|
pub token_budget: u32,
|
||||||
/// 12 小时空闲阈值(秒)
|
/// 12 小时空闲阈值(秒)
|
||||||
pub idle_timeout_secs: i64,
|
pub idle_timeout_secs: i64,
|
||||||
|
/// 会话 ID(用于追踪工具调用往返)
|
||||||
|
pub session_id: Uuid,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Default for ChatSession {
|
impl Default for ChatSession {
|
||||||
@@ -58,6 +61,7 @@ impl Default for ChatSession {
|
|||||||
last_user_at: None,
|
last_user_at: None,
|
||||||
token_budget: 28000,
|
token_budget: 28000,
|
||||||
idle_timeout_secs: 12 * 3600,
|
idle_timeout_secs: 12 * 3600,
|
||||||
|
session_id: Uuid::new_v4(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -90,6 +94,11 @@ impl ChatSession {
|
|||||||
.push(Message::tool_result(tool_call_id, name, result));
|
.push(Message::tool_result(tool_call_id, name, result));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// 添加一条系统说明消息(如已批准的工具)
|
||||||
|
pub fn add_system_note(&mut self, note: &str) {
|
||||||
|
self.messages.push(Message::system(note));
|
||||||
|
}
|
||||||
|
|
||||||
/// 是否因空闲超过阈值需要摘要
|
/// 是否因空闲超过阈值需要摘要
|
||||||
pub fn is_idle_timeout(&self) -> bool {
|
pub fn is_idle_timeout(&self) -> bool {
|
||||||
if let Some(last) = self.last_user_at {
|
if let Some(last) = self.last_user_at {
|
||||||
|
|||||||
+491
-716
File diff suppressed because it is too large
Load Diff
@@ -149,6 +149,7 @@ pub async fn recv_task(stream: &mut UnixStream) -> Result<TaskFrame, String> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// 接收并解析为 OutputFrame
|
/// 接收并解析为 OutputFrame
|
||||||
|
#[allow(dead_code)]
|
||||||
pub async fn recv_output(stream: &mut UnixStream) -> Result<OutputFrame, String> {
|
pub async fn recv_output(stream: &mut UnixStream) -> Result<OutputFrame, String> {
|
||||||
let v = recv_frame(stream).await?;
|
let v = recv_frame(stream).await?;
|
||||||
serde_json::from_value(v).map_err(|e| format!("解析 OutputFrame 失败: {}", e))
|
serde_json::from_value(v).map_err(|e| format!("解析 OutputFrame 失败: {}", e))
|
||||||
|
|||||||
+58
-5
@@ -4,6 +4,7 @@ use crate::llm::provider::{BoxedProvider, StreamReceiver, create_provider};
|
|||||||
use crate::llm::types::{ConversationConfig, Message, StreamChunk, ToolCall, Usage};
|
use crate::llm::types::{ConversationConfig, Message, StreamChunk, ToolCall, Usage};
|
||||||
use std::future::Future;
|
use std::future::Future;
|
||||||
use std::pin::Pin;
|
use std::pin::Pin;
|
||||||
|
use std::sync::atomic::{AtomicBool, Ordering};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use tokio::sync::Mutex;
|
use tokio::sync::Mutex;
|
||||||
|
|
||||||
@@ -18,11 +19,26 @@ pub type Summarizer = Arc<
|
|||||||
dyn Fn(String) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send>> + Send + Sync,
|
dyn Fn(String) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send>> + Send + Sync,
|
||||||
>;
|
>;
|
||||||
|
|
||||||
|
/// `chat_with_tools` 的返回结果
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct ChatResult {
|
||||||
|
pub reply: String,
|
||||||
|
pub used_tools: bool,
|
||||||
|
pub usage: Option<Usage>,
|
||||||
|
/// 是否有异步工具调用仍在等待结果
|
||||||
|
pub has_pending_async: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 异步工具标记 —— 工具执行器返回此前缀表示工具已异步入队
|
||||||
|
pub const ASYNC_MARKER: &str = "__ASYNC_PENDING__";
|
||||||
|
|
||||||
pub struct Conversation {
|
pub struct Conversation {
|
||||||
config: ConversationConfig,
|
config: ConversationConfig,
|
||||||
provider: BoxedProvider,
|
provider: BoxedProvider,
|
||||||
session: Arc<Mutex<ChatSession>>,
|
session: Arc<Mutex<ChatSession>>,
|
||||||
tool_executor: Option<ToolExecutor>,
|
tool_executor: Option<ToolExecutor>,
|
||||||
|
/// 是否有异步工具调用等待结果返回
|
||||||
|
pending_async_tools: Arc<AtomicBool>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Conversation {
|
impl Conversation {
|
||||||
@@ -33,12 +49,18 @@ impl Conversation {
|
|||||||
provider,
|
provider,
|
||||||
session: Arc::new(Mutex::new(ChatSession::new())),
|
session: Arc::new(Mutex::new(ChatSession::new())),
|
||||||
tool_executor: None,
|
tool_executor: None,
|
||||||
|
pending_async_tools: Arc::new(AtomicBool::new(false)),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn set_tool_executor(&mut self, executor: ToolExecutor) {
|
pub fn set_tool_executor(&mut self, executor: ToolExecutor) {
|
||||||
self.tool_executor = Some(executor);
|
self.tool_executor = Some(executor);
|
||||||
}
|
}
|
||||||
|
/// 是否有异步工具调用正在等待结果
|
||||||
|
#[allow(dead_code)]
|
||||||
|
pub fn has_pending_async(&self) -> bool {
|
||||||
|
self.pending_async_tools.load(Ordering::SeqCst)
|
||||||
|
}
|
||||||
pub fn session(&self) -> Arc<Mutex<ChatSession>> {
|
pub fn session(&self) -> Arc<Mutex<ChatSession>> {
|
||||||
Arc::clone(&self.session)
|
Arc::clone(&self.session)
|
||||||
}
|
}
|
||||||
@@ -113,10 +135,7 @@ impl Conversation {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// 对话 + 工具循环(带上下文管理)
|
/// 对话 + 工具循环(带上下文管理)
|
||||||
pub async fn chat_with_tools(
|
pub async fn chat_with_tools(&self, user_message: String) -> Result<ChatResult, String> {
|
||||||
&self,
|
|
||||||
user_message: String,
|
|
||||||
) -> Result<(String, bool, Option<Usage>), String> {
|
|
||||||
// ── 空闲超时检查 ──
|
// ── 空闲超时检查 ──
|
||||||
{
|
{
|
||||||
let s = self.session.lock().await;
|
let s = self.session.lock().await;
|
||||||
@@ -128,7 +147,17 @@ impl Conversation {
|
|||||||
}
|
}
|
||||||
|
|
||||||
self.session.lock().await.add_user(user_message.clone());
|
self.session.lock().await.add_user(user_message.clone());
|
||||||
|
self.run_tool_loop().await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 在异步工具结果返回后恢复工具循环(不追加用户消息)
|
||||||
|
pub async fn resume_tool_loop(&self) -> Result<ChatResult, String> {
|
||||||
|
self.pending_async_tools.store(false, Ordering::SeqCst);
|
||||||
|
self.run_tool_loop().await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 工具循环核心:调 LLM → 处理工具调用 → 循环直到无工具或需要异步等待
|
||||||
|
async fn run_tool_loop(&self) -> Result<ChatResult, String> {
|
||||||
let mut used_tools = false;
|
let mut used_tools = false;
|
||||||
let mut last_usage = None;
|
let mut last_usage = None;
|
||||||
let mut turn_count = 0u32;
|
let mut turn_count = 0u32;
|
||||||
@@ -185,6 +214,7 @@ impl Conversation {
|
|||||||
.await
|
.await
|
||||||
.add_assistant_tool_calls(calls.clone());
|
.add_assistant_tool_calls(calls.clone());
|
||||||
|
|
||||||
|
let mut has_async = false;
|
||||||
for tc in &calls {
|
for tc in &calls {
|
||||||
let result = match &self.tool_executor {
|
let result = match &self.tool_executor {
|
||||||
Some(exec) => {
|
Some(exec) => {
|
||||||
@@ -195,6 +225,14 @@ impl Conversation {
|
|||||||
}
|
}
|
||||||
None => "工具执行器未配置".to_string(),
|
None => "工具执行器未配置".to_string(),
|
||||||
};
|
};
|
||||||
|
|
||||||
|
if result.starts_with(ASYNC_MARKER) {
|
||||||
|
has_async = true;
|
||||||
|
tracing::info!("📦 {} → 异步入队", tc.function.name);
|
||||||
|
// 不添加到 session,等待 ToolResult 异步回喂
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
tracing::info!("📦 {} → {:.150}", tc.function.name, result);
|
tracing::info!("📦 {} → {:.150}", tc.function.name, result);
|
||||||
self.session.lock().await.add_tool_result(
|
self.session.lock().await.add_tool_result(
|
||||||
&tc.id,
|
&tc.id,
|
||||||
@@ -202,6 +240,16 @@ impl Conversation {
|
|||||||
&result,
|
&result,
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if has_async {
|
||||||
|
self.pending_async_tools.store(true, Ordering::SeqCst);
|
||||||
|
return Ok(ChatResult {
|
||||||
|
reply: full_text,
|
||||||
|
used_tools,
|
||||||
|
usage: last_usage,
|
||||||
|
has_pending_async: true,
|
||||||
|
});
|
||||||
|
}
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -224,7 +272,12 @@ impl Conversation {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
return Ok((full_text, used_tools, last_usage));
|
return Ok(ChatResult {
|
||||||
|
reply: full_text,
|
||||||
|
used_tools,
|
||||||
|
usage: last_usage,
|
||||||
|
has_pending_async: false,
|
||||||
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+1
-1
@@ -3,5 +3,5 @@ pub mod deepseek;
|
|||||||
pub mod provider;
|
pub mod provider;
|
||||||
pub mod types;
|
pub mod types;
|
||||||
|
|
||||||
pub use conversation::{Conversation, ToolExecutor, DEFAULT_SYSTEM_PROMPT};
|
pub use conversation::{ChatResult, Conversation, ToolExecutor, ASYNC_MARKER, DEFAULT_SYSTEM_PROMPT};
|
||||||
pub use types::{ConversationConfig, Usage};
|
pub use types::{ConversationConfig, Usage};
|
||||||
|
|||||||
+7
-3
@@ -1,3 +1,4 @@
|
|||||||
|
mod channel;
|
||||||
mod cli;
|
mod cli;
|
||||||
mod context;
|
mod context;
|
||||||
mod daemon;
|
mod daemon;
|
||||||
@@ -5,6 +6,7 @@ mod db;
|
|||||||
mod ipc;
|
mod ipc;
|
||||||
mod llm;
|
mod llm;
|
||||||
mod logger;
|
mod logger;
|
||||||
|
mod queue;
|
||||||
mod scheduler;
|
mod scheduler;
|
||||||
mod state;
|
mod state;
|
||||||
mod tools;
|
mod tools;
|
||||||
@@ -15,11 +17,11 @@ use clap::Parser;
|
|||||||
use cli::{AmapAction, Cli, Commands, MemosAction, ToolCommand};
|
use cli::{AmapAction, Cli, Commands, MemosAction, ToolCommand};
|
||||||
use context::MemoryStore;
|
use context::MemoryStore;
|
||||||
use db::Database;
|
use db::Database;
|
||||||
use llm::{Conversation, ConversationConfig, DEFAULT_SYSTEM_PROMPT, ToolExecutor, Usage};
|
use llm::{ChatResult, Conversation, ConversationConfig, DEFAULT_SYSTEM_PROMPT, ToolExecutor, Usage};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use tools::approval::{ApprovalDecision, ApprovalManager};
|
use tools::approval::{ApprovalDecision, ApprovalManager};
|
||||||
use tools::types::ExecutionContext;
|
use tools::types::ExecutionContext;
|
||||||
use tracing::{error, info};
|
use tracing::{error, info, warn};
|
||||||
use wechat::client::WeChatClient;
|
use wechat::client::WeChatClient;
|
||||||
|
|
||||||
#[tokio::main]
|
#[tokio::main]
|
||||||
@@ -105,6 +107,8 @@ async fn main() {
|
|||||||
cmd_daemon(sock, &database, &file_state, &memory_store).await;
|
cmd_daemon(sock, &database, &file_state, &memory_store).await;
|
||||||
}
|
}
|
||||||
Commands::Worker { sock } => {
|
Commands::Worker { sock } => {
|
||||||
|
warn!("Worker 模式已废弃,请使用 daemon 模式");
|
||||||
|
#[allow(deprecated)]
|
||||||
if let Err(e) = worker::run(&sock).await {
|
if let Err(e) = worker::run(&sock).await {
|
||||||
error!("Worker 失败: {}", e);
|
error!("Worker 失败: {}", e);
|
||||||
}
|
}
|
||||||
@@ -665,7 +669,7 @@ async fn generate_reply_with_tools(
|
|||||||
user_text: String,
|
user_text: String,
|
||||||
db: &Option<Arc<Database>>,
|
db: &Option<Arc<Database>>,
|
||||||
) -> Result<String, String> {
|
) -> Result<String, String> {
|
||||||
let (reply, _used_tools, usage) = conv.chat_with_tools(user_text).await?;
|
let ChatResult { reply, used_tools: _used_tools, usage, .. } = conv.chat_with_tools(user_text).await?;
|
||||||
if let Some(u) = &usage {
|
if let Some(u) = &usage {
|
||||||
info!(
|
info!(
|
||||||
"📊 Token: {} 总 / {} 提示 / {} 生成 | 缓存命中 {} | 缓存未命中 {}",
|
"📊 Token: {} 总 / {} 提示 / {} 生成 | 缓存命中 {} | 缓存未命中 {}",
|
||||||
|
|||||||
@@ -0,0 +1,552 @@
|
|||||||
|
//! ## 多渠道消息队列 — MessageQueue
|
||||||
|
//!
|
||||||
|
//! # 设计目标
|
||||||
|
//!
|
||||||
|
//! 多用户场景下,保证公平性:
|
||||||
|
//! - A 发了 3 条消息、B 发了 1 条 → 处理顺序 A1, B1, A2, A3
|
||||||
|
//! - 不会因为 A 消息多而饿死 B
|
||||||
|
//!
|
||||||
|
//! 每个消息携带 user_id + correlation_id,保证回复能追溯到正确的用户。
|
||||||
|
//!
|
||||||
|
//! # 架构位置
|
||||||
|
//!
|
||||||
|
//! ```text
|
||||||
|
//! receive_messages → enqueue PipelineMessage
|
||||||
|
//! ┌─────────────────────────────┐
|
||||||
|
//! │ MessageQueue │
|
||||||
|
//! │ [User A] [User B] │ ← 每用户独立 VecDeque
|
||||||
|
//! │ pending: [A, B, A] │ ← 公平轮转索引
|
||||||
|
//! └─────────────────────────────┘
|
||||||
|
//! QueueRunner::run() → dequeue → 按 kind 路由到 mpsc 通道
|
||||||
|
//! ├─ llm_tx (UserMessage, ToolResult)
|
||||||
|
//! ├─ tool_tx (ToolCall, ApprovalRequest)
|
||||||
|
//! └─ send_tx (LLMReply)
|
||||||
|
//! ```
|
||||||
|
|
||||||
|
use crate::channel::ChannelId;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::collections::{HashMap, VecDeque};
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::time::Duration;
|
||||||
|
use tokio::sync::Notify;
|
||||||
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
// ─── 管线消息类型 ───
|
||||||
|
|
||||||
|
/// 消息种类 —— 区分不同消费路由
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub enum MessageKind {
|
||||||
|
/// 从外部收到的用户消息 → 交给 LLM Consumer
|
||||||
|
UserMessage {
|
||||||
|
text: String,
|
||||||
|
message_id: String,
|
||||||
|
context_token: Option<String>,
|
||||||
|
/// 审批通过后携带的已批准工具名
|
||||||
|
approved_tool: Option<String>,
|
||||||
|
},
|
||||||
|
/// LLM 生成的回复文本 → 发给 Send Consumer
|
||||||
|
LLMReply {
|
||||||
|
text: String,
|
||||||
|
context_token: Option<String>,
|
||||||
|
},
|
||||||
|
/// LLM 请求的工具调用 → 发给 Tool Consumer
|
||||||
|
ToolCall {
|
||||||
|
session_id: Uuid,
|
||||||
|
tool_call_id: String,
|
||||||
|
tool_name: String,
|
||||||
|
arguments: String,
|
||||||
|
context_token: Option<String>,
|
||||||
|
},
|
||||||
|
/// 工具执行结果 → 回喂 LLM Consumer
|
||||||
|
ToolResult {
|
||||||
|
session_id: Uuid,
|
||||||
|
tool_call_id: String,
|
||||||
|
tool_name: String,
|
||||||
|
result: String,
|
||||||
|
context_token: Option<String>,
|
||||||
|
},
|
||||||
|
/// 高风险工具 → 需用户审批
|
||||||
|
ApprovalRequest {
|
||||||
|
session_id: Uuid,
|
||||||
|
tool_name: String,
|
||||||
|
tool_args: String,
|
||||||
|
/// 审批后要追加到用户消息中的文本
|
||||||
|
approval_prompt: String,
|
||||||
|
},
|
||||||
|
/// 定时任务到期 → 交给 LLM Consumer
|
||||||
|
ScheduledTask {
|
||||||
|
text: String,
|
||||||
|
task_name: String,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 管线消息 —— 携带路由元数据的统一消息类型
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct PipelineMessage {
|
||||||
|
/// 来源/目标渠道(平台 + 用户),全程不变
|
||||||
|
pub channel: ChannelId,
|
||||||
|
/// 关联 ID:追踪一次用户请求的完整链路(用户消息→LLM→工具→回喂→发送)
|
||||||
|
pub correlation_id: Uuid,
|
||||||
|
/// 消息种类
|
||||||
|
pub kind: MessageKind,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PipelineMessage {
|
||||||
|
/// 便捷:渠道中的 user_id
|
||||||
|
pub fn user_id(&self) -> &str { &self.channel.user_id }
|
||||||
|
/// 便捷:渠道中的 platform
|
||||||
|
#[allow(dead_code)]
|
||||||
|
pub fn platform(&self) -> &str { &self.channel.platform }
|
||||||
|
/// 快速构造一条用户消息
|
||||||
|
#[allow(dead_code)]
|
||||||
|
pub fn user_message(
|
||||||
|
user_id: impl Into<String>,
|
||||||
|
text: impl Into<String>,
|
||||||
|
message_id: impl Into<String>,
|
||||||
|
context_token: Option<String>,
|
||||||
|
) -> Self {
|
||||||
|
Self {
|
||||||
|
channel: ChannelId::wechat(user_id),
|
||||||
|
correlation_id: Uuid::new_v4(),
|
||||||
|
kind: MessageKind::UserMessage {
|
||||||
|
text: text.into(),
|
||||||
|
message_id: message_id.into(),
|
||||||
|
context_token,
|
||||||
|
approved_tool: None,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 快速构造一条 LLM 回复
|
||||||
|
pub fn llm_reply(
|
||||||
|
channel: ChannelId,
|
||||||
|
correlation_id: Uuid,
|
||||||
|
text: impl Into<String>,
|
||||||
|
context_token: Option<String>,
|
||||||
|
) -> Self {
|
||||||
|
Self {
|
||||||
|
channel,
|
||||||
|
correlation_id,
|
||||||
|
kind: MessageKind::LLMReply {
|
||||||
|
text: text.into(),
|
||||||
|
context_token,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 快速构造一条工具调用
|
||||||
|
pub fn tool_call(
|
||||||
|
channel: ChannelId,
|
||||||
|
correlation_id: Uuid,
|
||||||
|
session_id: Uuid,
|
||||||
|
tool_call_id: impl Into<String>,
|
||||||
|
tool_name: impl Into<String>,
|
||||||
|
arguments: impl Into<String>,
|
||||||
|
context_token: Option<String>,
|
||||||
|
) -> Self {
|
||||||
|
Self {
|
||||||
|
channel,
|
||||||
|
correlation_id,
|
||||||
|
kind: MessageKind::ToolCall {
|
||||||
|
session_id,
|
||||||
|
tool_call_id: tool_call_id.into(),
|
||||||
|
tool_name: tool_name.into(),
|
||||||
|
arguments: arguments.into(),
|
||||||
|
context_token,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 快速构造一条工具结果
|
||||||
|
pub fn tool_result(
|
||||||
|
channel: ChannelId,
|
||||||
|
correlation_id: Uuid,
|
||||||
|
session_id: Uuid,
|
||||||
|
tool_call_id: impl Into<String>,
|
||||||
|
tool_name: impl Into<String>,
|
||||||
|
result: impl Into<String>,
|
||||||
|
context_token: Option<String>,
|
||||||
|
) -> Self {
|
||||||
|
Self {
|
||||||
|
channel,
|
||||||
|
correlation_id,
|
||||||
|
kind: MessageKind::ToolResult {
|
||||||
|
session_id,
|
||||||
|
tool_call_id: tool_call_id.into(),
|
||||||
|
tool_name: tool_name.into(),
|
||||||
|
result: result.into(),
|
||||||
|
context_token,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── 公平轮转队列 ───
|
||||||
|
|
||||||
|
/// 多渠道公平轮转消息队列
|
||||||
|
///
|
||||||
|
/// 按 ChannelId(platform + user_id)分队列,轮转出队保证公平。
|
||||||
|
/// 出队后由 QueueRunner 按 MessageKind 路由到对应 mpsc 通道。
|
||||||
|
pub struct MessageQueue {
|
||||||
|
/// 每渠道独立队列
|
||||||
|
queues: HashMap<ChannelId, VecDeque<PipelineMessage>>,
|
||||||
|
/// 有待处理消息的渠道(轮转顺序)
|
||||||
|
pending: VecDeque<ChannelId>,
|
||||||
|
/// 空→非空时触发 Notify 唤醒消费者
|
||||||
|
signal: Arc<Notify>,
|
||||||
|
/// 定时唤醒间隔(兜底用)
|
||||||
|
poll_interval: Duration,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl MessageQueue {
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
queues: HashMap::new(),
|
||||||
|
pending: VecDeque::new(),
|
||||||
|
signal: Arc::new(Notify::new()),
|
||||||
|
poll_interval: Duration::from_millis(200),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[allow(dead_code)]
|
||||||
|
pub fn with_poll_interval(mut self, d: Duration) -> Self {
|
||||||
|
self.poll_interval = d;
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 入队。如果该渠道之前为空 → 加入 pending 尾部 + 通知
|
||||||
|
pub fn enqueue(&mut self, item: PipelineMessage) {
|
||||||
|
let ch = item.channel.clone();
|
||||||
|
let q = self.queues.entry(ch.clone()).or_default();
|
||||||
|
let was_empty = q.is_empty();
|
||||||
|
q.push_back(item);
|
||||||
|
if was_empty {
|
||||||
|
self.pending.push_back(ch);
|
||||||
|
self.signal.notify_one();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 公平轮转出队
|
||||||
|
///
|
||||||
|
/// 每次调用切换到下一个有消息的渠道:
|
||||||
|
/// - 从 pending 头部取出一个渠道
|
||||||
|
/// - pop 其一条消息
|
||||||
|
/// - 该渠道还有消息 → 放回 pending 尾部
|
||||||
|
pub fn dequeue(&mut self) -> Option<PipelineMessage> {
|
||||||
|
while let Some(ch) = self.pending.pop_front() {
|
||||||
|
if let Some(q) = self.queues.get_mut(&ch) {
|
||||||
|
if let Some(item) = q.pop_front() {
|
||||||
|
if !q.is_empty() {
|
||||||
|
self.pending.push_back(ch);
|
||||||
|
} else {
|
||||||
|
self.queues.remove(&ch);
|
||||||
|
}
|
||||||
|
return Some(item);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
self.queues.remove(&ch);
|
||||||
|
}
|
||||||
|
None
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn signal(&self) -> Arc<Notify> {
|
||||||
|
Arc::clone(&self.signal)
|
||||||
|
}
|
||||||
|
#[allow(dead_code)]
|
||||||
|
pub fn poll_interval(&self) -> Duration {
|
||||||
|
self.poll_interval
|
||||||
|
}
|
||||||
|
|
||||||
|
#[allow(dead_code)]
|
||||||
|
pub fn has_pending(&self, ch: &ChannelId) -> bool {
|
||||||
|
self.queues.get(ch).map_or(false, |q| !q.is_empty())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[allow(dead_code)]
|
||||||
|
pub fn pending_count(&self) -> usize {
|
||||||
|
self.pending.len()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 清空指定渠道的全部待处理消息
|
||||||
|
#[allow(dead_code)]
|
||||||
|
pub fn drain_channel(&mut self, ch: &ChannelId) -> Vec<PipelineMessage> {
|
||||||
|
let drained = self
|
||||||
|
.queues
|
||||||
|
.remove(ch)
|
||||||
|
.map(|mut q| q.drain(..).collect())
|
||||||
|
.unwrap_or_default();
|
||||||
|
self.pending.retain(|c| c != ch);
|
||||||
|
drained
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── 消费者通道标识 ───
|
||||||
|
|
||||||
|
/// 消息路由目标
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
pub enum ConsumerTarget {
|
||||||
|
/// LLM 消费者(处理 UserMessage, ToolResult)
|
||||||
|
LLM,
|
||||||
|
/// 工具消费者(处理 ToolCall, ApprovalRequest)
|
||||||
|
Tool,
|
||||||
|
/// 发送消费者(处理 LLMReply)
|
||||||
|
Send,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PipelineMessage {
|
||||||
|
/// 判断这条消息应该路由到哪个消费者
|
||||||
|
pub fn target(&self) -> ConsumerTarget {
|
||||||
|
match self.kind {
|
||||||
|
MessageKind::UserMessage { .. }
|
||||||
|
| MessageKind::ToolResult { .. }
|
||||||
|
| MessageKind::ScheduledTask { .. } => ConsumerTarget::LLM,
|
||||||
|
MessageKind::ToolCall { .. } | MessageKind::ApprovalRequest { .. } => {
|
||||||
|
ConsumerTarget::Tool
|
||||||
|
}
|
||||||
|
MessageKind::LLMReply { .. } => ConsumerTarget::Send,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── 测试 ───
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use crate::queue::QueueRunner;
|
||||||
|
|
||||||
|
fn make_user_msg(user: &str, text: &str) -> PipelineMessage {
|
||||||
|
PipelineMessage::user_message(user, text, "", None)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_target_routing() {
|
||||||
|
let msg = PipelineMessage::user_message("u1", "hello", "m1", None);
|
||||||
|
assert_eq!(msg.target(), ConsumerTarget::LLM);
|
||||||
|
|
||||||
|
let msg = PipelineMessage::llm_reply(ChannelId::wechat("u1"), Uuid::new_v4(), "hi", None);
|
||||||
|
assert_eq!(msg.target(), ConsumerTarget::Send);
|
||||||
|
|
||||||
|
let msg = PipelineMessage::tool_call(
|
||||||
|
ChannelId::wechat("u1"),
|
||||||
|
Uuid::new_v4(),
|
||||||
|
Uuid::new_v4(),
|
||||||
|
"tc1",
|
||||||
|
"query_weather",
|
||||||
|
r#"{"location":"北京"}"#,
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
assert_eq!(msg.target(), ConsumerTarget::Tool);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_fair_round_robin() {
|
||||||
|
let mut q = MessageQueue::new();
|
||||||
|
q.enqueue(make_user_msg("A", "a1"));
|
||||||
|
q.enqueue(make_user_msg("A", "a2"));
|
||||||
|
q.enqueue(make_user_msg("A", "a3"));
|
||||||
|
q.enqueue(make_user_msg("B", "b1"));
|
||||||
|
assert_eq!(q.dequeue().unwrap().channel.user_id, "A");
|
||||||
|
assert_eq!(q.dequeue().unwrap().channel.user_id, "B");
|
||||||
|
assert_eq!(q.dequeue().unwrap().channel.user_id, "A");
|
||||||
|
assert_eq!(q.dequeue().unwrap().channel.user_id, "A");
|
||||||
|
assert!(q.dequeue().is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_single_user() {
|
||||||
|
let mut q = MessageQueue::new();
|
||||||
|
q.enqueue(make_user_msg("X", "hi"));
|
||||||
|
q.enqueue(make_user_msg("X", "hello"));
|
||||||
|
assert_eq!(q.dequeue().unwrap().channel.user_id, "X");
|
||||||
|
assert_eq!(q.dequeue().unwrap().channel.user_id, "X");
|
||||||
|
assert!(q.dequeue().is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_empty_queue() {
|
||||||
|
let mut q = MessageQueue::new();
|
||||||
|
assert!(q.dequeue().is_none());
|
||||||
|
assert_eq!(q.pending_count(), 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_user_id_preserved() {
|
||||||
|
let msg = PipelineMessage::user_message("wx_user_1", "天气?", "m1", None);
|
||||||
|
assert_eq!(msg.channel.user_id, "wx_user_1");
|
||||||
|
|
||||||
|
let reply = PipelineMessage::llm_reply(ChannelId::wechat("wx_user_1"), msg.correlation_id, "晴天", None);
|
||||||
|
assert_eq!(reply.channel.user_id, "wx_user_1");
|
||||||
|
assert_eq!(reply.correlation_id, msg.correlation_id);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_mixed_item_types() {
|
||||||
|
let mut q = MessageQueue::new();
|
||||||
|
q.enqueue(PipelineMessage::user_message("U", "user msg", "1", None));
|
||||||
|
q.enqueue(PipelineMessage {
|
||||||
|
channel: ChannelId::wechat("U"),
|
||||||
|
correlation_id: Uuid::new_v4(),
|
||||||
|
kind: MessageKind::ScheduledTask {
|
||||||
|
text: "task result".into(),
|
||||||
|
task_name: "reminder".into(),
|
||||||
|
},
|
||||||
|
});
|
||||||
|
assert!(matches!(
|
||||||
|
q.dequeue().unwrap().kind,
|
||||||
|
MessageKind::UserMessage { .. }
|
||||||
|
));
|
||||||
|
assert!(matches!(
|
||||||
|
q.dequeue().unwrap().kind,
|
||||||
|
MessageKind::ScheduledTask { .. }
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── 端到端关联测试 ───
|
||||||
|
|
||||||
|
/// 验证 user_id 从用户消息到 LLM 回复全程不变
|
||||||
|
#[test]
|
||||||
|
fn test_user_id_pipeline_preserved() {
|
||||||
|
let user_id = "wx_user_1";
|
||||||
|
|
||||||
|
// 用户消息入队
|
||||||
|
let user_msg = PipelineMessage::user_message(user_id, "天气怎么样?", "msg1", None);
|
||||||
|
assert_eq!(user_msg.channel.user_id, user_id);
|
||||||
|
let correlation_id = user_msg.correlation_id;
|
||||||
|
|
||||||
|
// LLM 回复(模拟 LLM Consumer 生成)
|
||||||
|
let reply = PipelineMessage::llm_reply(ChannelId::wechat(user_id), correlation_id, "今天晴天", None);
|
||||||
|
assert_eq!(reply.channel.user_id, user_id);
|
||||||
|
assert_eq!(reply.correlation_id, correlation_id);
|
||||||
|
|
||||||
|
// 工具调用(模拟 LLM Consumer 生成)
|
||||||
|
let tool_call = PipelineMessage::tool_call(
|
||||||
|
ChannelId::wechat(user_id),
|
||||||
|
correlation_id,
|
||||||
|
Uuid::new_v4(),
|
||||||
|
"tc1",
|
||||||
|
"query_weather",
|
||||||
|
r#"{"location":"北京"}"#,
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
assert_eq!(tool_call.channel.user_id, user_id);
|
||||||
|
assert_eq!(tool_call.correlation_id, correlation_id);
|
||||||
|
|
||||||
|
// 工具结果回喂(模拟 Tool Consumer 生成)
|
||||||
|
let tool_result = PipelineMessage::tool_result(
|
||||||
|
ChannelId::wechat(user_id),
|
||||||
|
correlation_id,
|
||||||
|
Uuid::new_v4(),
|
||||||
|
"tc1",
|
||||||
|
"query_weather",
|
||||||
|
"25°C",
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
assert_eq!(tool_result.channel.user_id, user_id);
|
||||||
|
assert_eq!(tool_result.correlation_id, correlation_id);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 验证多个用户的消息不会串
|
||||||
|
#[test]
|
||||||
|
fn test_multi_user_isolation() {
|
||||||
|
let mut q = MessageQueue::new();
|
||||||
|
|
||||||
|
q.enqueue(PipelineMessage::user_message("u1", "你好", "1", None));
|
||||||
|
q.enqueue(PipelineMessage::user_message("u2", "hello", "2", None));
|
||||||
|
q.enqueue(PipelineMessage::user_message("u1", "天气?", "3", None));
|
||||||
|
|
||||||
|
// 按公平轮转出队,验证 user_id 正确
|
||||||
|
let msg1 = q.dequeue().unwrap();
|
||||||
|
assert_eq!(msg1.channel.user_id, "u1");
|
||||||
|
let msg2 = q.dequeue().unwrap();
|
||||||
|
assert_eq!(msg2.channel.user_id, "u2");
|
||||||
|
let msg3 = q.dequeue().unwrap();
|
||||||
|
assert_eq!(msg3.channel.user_id, "u1");
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 验证路由目标正确
|
||||||
|
#[test]
|
||||||
|
fn test_consumer_target_routing() {
|
||||||
|
// UserMessage → LLM
|
||||||
|
let m = PipelineMessage::user_message("u1", "hi", "1", None);
|
||||||
|
assert_eq!(m.target(), ConsumerTarget::LLM);
|
||||||
|
|
||||||
|
// ToolResult → LLM
|
||||||
|
let m = PipelineMessage::tool_result(ChannelId::wechat("u1"), Uuid::new_v4(), Uuid::new_v4(), "tc1", "test", "ok", None);
|
||||||
|
assert_eq!(m.target(), ConsumerTarget::LLM);
|
||||||
|
|
||||||
|
// ScheduledTask → LLM
|
||||||
|
let m = PipelineMessage {
|
||||||
|
channel: ChannelId::wechat("u1"),
|
||||||
|
correlation_id: Uuid::new_v4(),
|
||||||
|
kind: MessageKind::ScheduledTask { text: "task".into(), task_name: "t1".into() },
|
||||||
|
};
|
||||||
|
assert_eq!(m.target(), ConsumerTarget::LLM);
|
||||||
|
|
||||||
|
// ToolCall → Tool
|
||||||
|
let m = PipelineMessage::tool_call(ChannelId::wechat("u1"), Uuid::new_v4(), Uuid::new_v4(), "tc1", "weather", "{}", None);
|
||||||
|
assert_eq!(m.target(), ConsumerTarget::Tool);
|
||||||
|
|
||||||
|
// ApprovalRequest → Tool
|
||||||
|
let m = PipelineMessage {
|
||||||
|
channel: ChannelId::wechat("u1"),
|
||||||
|
correlation_id: Uuid::new_v4(),
|
||||||
|
kind: MessageKind::ApprovalRequest {
|
||||||
|
session_id: Uuid::new_v4(),
|
||||||
|
tool_name: "high_risk_tool".into(),
|
||||||
|
tool_args: "{}".into(),
|
||||||
|
approval_prompt: "".into(),
|
||||||
|
},
|
||||||
|
};
|
||||||
|
assert_eq!(m.target(), ConsumerTarget::Tool);
|
||||||
|
|
||||||
|
// LLMReply → Send
|
||||||
|
let m = PipelineMessage::llm_reply(ChannelId::wechat("u1"), Uuid::new_v4(), "reply", None);
|
||||||
|
assert_eq!(m.target(), ConsumerTarget::Send);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 验证 QueueRunner 的路由逻辑:入队 → 路由到正确通道
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_queue_runner_routing() {
|
||||||
|
let (runner, mut channels) = QueueRunner::new(16);
|
||||||
|
let enqueue = runner.enqueue_handle();
|
||||||
|
|
||||||
|
// 启动 runner(在后台运行)
|
||||||
|
let handle = tokio::spawn(async move {
|
||||||
|
runner.run().await;
|
||||||
|
});
|
||||||
|
|
||||||
|
// 入队一个 UserMessage → 应到 llm_rx
|
||||||
|
let user_msg = PipelineMessage::user_message("u1", "hello", "m1", None);
|
||||||
|
enqueue.enqueue(user_msg).await;
|
||||||
|
|
||||||
|
// 入队一个 LLMReply → 应到 send_rx
|
||||||
|
let reply_msg = PipelineMessage::llm_reply(ChannelId::wechat("u1"), Uuid::new_v4(), "reply", None);
|
||||||
|
enqueue.enqueue(reply_msg).await;
|
||||||
|
|
||||||
|
// 验证 llm_rx 收到 UserMessage
|
||||||
|
let llm_msg = tokio::time::timeout(
|
||||||
|
std::time::Duration::from_secs(2),
|
||||||
|
channels.llm_rx.recv(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect("llm_rx 超时")
|
||||||
|
.expect("llm_rx 通道关闭");
|
||||||
|
assert_eq!(llm_msg.channel.user_id, "u1");
|
||||||
|
assert!(matches!(llm_msg.kind, MessageKind::UserMessage { .. }));
|
||||||
|
|
||||||
|
// 验证 send_rx 收到 LLMReply
|
||||||
|
let send_msg = tokio::time::timeout(
|
||||||
|
std::time::Duration::from_secs(2),
|
||||||
|
channels.send_rx.recv(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect("send_rx 超时")
|
||||||
|
.expect("send_rx 通道关闭");
|
||||||
|
assert_eq!(send_msg.channel.user_id, "u1");
|
||||||
|
assert!(matches!(send_msg.kind, MessageKind::LLMReply { .. }));
|
||||||
|
|
||||||
|
// 关闭 runner
|
||||||
|
handle.abort();
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,5 @@
|
|||||||
|
pub mod message_queue;
|
||||||
|
pub mod runner;
|
||||||
|
|
||||||
|
pub use message_queue::{ConsumerTarget, MessageKind, MessageQueue, PipelineMessage};
|
||||||
|
pub use runner::{ConsumerChannels, EnqueueHandle, QueueRunner};
|
||||||
@@ -0,0 +1,212 @@
|
|||||||
|
//! QueueRunner — 消息队列运行时
|
||||||
|
//!
|
||||||
|
//! 从公平轮转 MessageQueue 中出队消息,按 MessageKind 路由到对应的
|
||||||
|
//! tokio mpsc 通道,由注册的消费者 task 处理。
|
||||||
|
//!
|
||||||
|
//! # 架构
|
||||||
|
//!
|
||||||
|
//! ```text
|
||||||
|
//! ┌──────────────┐
|
||||||
|
//! enqueue ──────→ │ MessageQueue │ ← 公平轮转
|
||||||
|
//! └──────┬───────┘
|
||||||
|
//! │ dequeue
|
||||||
|
//! ▼
|
||||||
|
//! ┌──────────────┐
|
||||||
|
//! │ QueueRunner │ ← 路由循环
|
||||||
|
//! └──┬───┬───┬──┘
|
||||||
|
//! ┌─┘ │ └─┐
|
||||||
|
//! ▼ ▼ ▼
|
||||||
|
//! llm_tx tool_tx send_tx
|
||||||
|
//! │ │ │
|
||||||
|
//! ┌────┘ ┌─┘ ┌──┘
|
||||||
|
//! ▼ ▼ ▼
|
||||||
|
//! LLMCons. ToolCons. SendCons.
|
||||||
|
//! ```
|
||||||
|
|
||||||
|
use crate::queue::{ConsumerTarget, MessageQueue, PipelineMessage};
|
||||||
|
use std::sync::atomic::{AtomicBool, Ordering};
|
||||||
|
use std::sync::Arc;
|
||||||
|
use tokio::sync::{mpsc, Mutex, Notify};
|
||||||
|
use tracing::{error, info, warn};
|
||||||
|
|
||||||
|
/// 消费者通道集 —— 创建 QueueRunner 时获得接收端
|
||||||
|
pub struct ConsumerChannels {
|
||||||
|
pub llm_rx: mpsc::Receiver<PipelineMessage>,
|
||||||
|
pub tool_rx: mpsc::Receiver<PipelineMessage>,
|
||||||
|
pub send_rx: mpsc::Receiver<PipelineMessage>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 队列运行时 —— 负责从公平轮转队列出队并按类型路由
|
||||||
|
pub struct QueueRunner {
|
||||||
|
/// 公平轮转队列
|
||||||
|
queue: Arc<Mutex<MessageQueue>>,
|
||||||
|
/// LLM 消费者通道
|
||||||
|
llm_tx: mpsc::Sender<PipelineMessage>,
|
||||||
|
/// 工具消费者通道
|
||||||
|
tool_tx: mpsc::Sender<PipelineMessage>,
|
||||||
|
/// 发送消费者通道
|
||||||
|
send_tx: mpsc::Sender<PipelineMessage>,
|
||||||
|
/// 关闭信号
|
||||||
|
shutdown: Arc<AtomicBool>,
|
||||||
|
/// 队列信号(空→非空通知)
|
||||||
|
signal: Arc<Notify>,
|
||||||
|
/// 轮询间隔
|
||||||
|
poll_interval: std::time::Duration,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl QueueRunner {
|
||||||
|
/// 创建 QueueRunner 和对应的消费者通道
|
||||||
|
///
|
||||||
|
/// `channel_size`: 每个 mpsc 通道的缓冲区大小
|
||||||
|
pub fn new(channel_size: usize) -> (Self, ConsumerChannels) {
|
||||||
|
let (llm_tx, llm_rx) = mpsc::channel(channel_size);
|
||||||
|
let (tool_tx, tool_rx) = mpsc::channel(channel_size);
|
||||||
|
let (send_tx, send_rx) = mpsc::channel(channel_size);
|
||||||
|
|
||||||
|
let queue = MessageQueue::new();
|
||||||
|
let signal = queue.signal();
|
||||||
|
let queue = Arc::new(Mutex::new(queue));
|
||||||
|
|
||||||
|
let runner = Self {
|
||||||
|
queue,
|
||||||
|
llm_tx,
|
||||||
|
tool_tx,
|
||||||
|
send_tx,
|
||||||
|
shutdown: Arc::new(AtomicBool::new(false)),
|
||||||
|
signal,
|
||||||
|
poll_interval: std::time::Duration::from_millis(200),
|
||||||
|
};
|
||||||
|
|
||||||
|
let channels = ConsumerChannels {
|
||||||
|
llm_rx,
|
||||||
|
tool_rx,
|
||||||
|
send_rx,
|
||||||
|
};
|
||||||
|
|
||||||
|
(runner, channels)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 获取队列的写句柄(用于外部入队)
|
||||||
|
pub fn enqueue_handle(&self) -> EnqueueHandle {
|
||||||
|
EnqueueHandle {
|
||||||
|
queue: self.queue.clone(),
|
||||||
|
shutdown: self.shutdown.clone(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 获取关闭信号句柄
|
||||||
|
pub fn shutdown_handle(&self) -> ShutdownHandle {
|
||||||
|
ShutdownHandle {
|
||||||
|
shutdown: self.shutdown.clone(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 启动路由主循环(阻塞,应在 tokio::spawn 中运行)
|
||||||
|
pub async fn run(&self) {
|
||||||
|
info!("QueueRunner 已启动");
|
||||||
|
loop {
|
||||||
|
// 检查关闭信号
|
||||||
|
if self.shutdown.load(Ordering::Relaxed) {
|
||||||
|
info!("QueueRunner 收到关闭信号,停止路由");
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
|
// 出队
|
||||||
|
let item = {
|
||||||
|
let mut q = self.queue.lock().await;
|
||||||
|
q.dequeue()
|
||||||
|
};
|
||||||
|
|
||||||
|
if let Some(msg) = item {
|
||||||
|
// 预提取字段(避免 borrow + move 冲突)
|
||||||
|
let target = msg.target();
|
||||||
|
let user_id = msg.channel.user_id.clone();
|
||||||
|
let correlation_id = msg.correlation_id;
|
||||||
|
|
||||||
|
let result = match target {
|
||||||
|
ConsumerTarget::LLM => self.llm_tx.send(msg).await,
|
||||||
|
ConsumerTarget::Tool => self.tool_tx.send(msg).await,
|
||||||
|
ConsumerTarget::Send => self.send_tx.send(msg).await,
|
||||||
|
};
|
||||||
|
|
||||||
|
if let Err(_e) = result {
|
||||||
|
warn!(
|
||||||
|
"路由消息失败 (target={:?}, user={}, corr={}): 通道已关闭,跳过",
|
||||||
|
target, user_id, correlation_id
|
||||||
|
);
|
||||||
|
// 不中断整个路由循环,跳过该消息继续处理
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
// 队列为空,等待通知或超时
|
||||||
|
tokio::select! {
|
||||||
|
_ = self.signal.notified() => {
|
||||||
|
// 有新的消息入队,继续循环
|
||||||
|
}
|
||||||
|
_ = tokio::time::sleep(self.poll_interval) => {
|
||||||
|
// 超时兜底,继续检查
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
info!("QueueRunner 已停止");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 入队句柄 —— 外部模块通过它向队列投放消息
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct EnqueueHandle {
|
||||||
|
queue: Arc<Mutex<MessageQueue>>,
|
||||||
|
shutdown: Arc<AtomicBool>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl EnqueueHandle {
|
||||||
|
/// 向队列投放一条消息
|
||||||
|
pub async fn enqueue(&self, msg: PipelineMessage) {
|
||||||
|
if self.shutdown.load(Ordering::Relaxed) {
|
||||||
|
warn!(
|
||||||
|
"队列已关闭,丢弃消息: user={}",
|
||||||
|
msg.user_id()
|
||||||
|
);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
self.queue.lock().await.enqueue(msg);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 关闭句柄
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct ShutdownHandle {
|
||||||
|
shutdown: Arc<AtomicBool>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ShutdownHandle {
|
||||||
|
/// 发送关闭信号
|
||||||
|
pub fn shutdown(&self) {
|
||||||
|
self.shutdown.store(true, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 等所有消费者处理完并关闭后,等待队列排空
|
||||||
|
#[allow(dead_code)]
|
||||||
|
pub async fn drain_and_shutdown(
|
||||||
|
_enqueue: &EnqueueHandle,
|
||||||
|
runner_shutdown: &ShutdownHandle,
|
||||||
|
consumer_handles: Vec<tokio::task::JoinHandle<()>>,
|
||||||
|
) {
|
||||||
|
info!("开始关闭队列系统...");
|
||||||
|
|
||||||
|
// 1. 阻止新消息入队(设置关闭标志)
|
||||||
|
runner_shutdown.shutdown();
|
||||||
|
|
||||||
|
// 2. 等待排空(给一个合理的时间)
|
||||||
|
tokio::time::sleep(std::time::Duration::from_millis(500)).await;
|
||||||
|
|
||||||
|
// 3. 等待消费者 task 完成
|
||||||
|
for handle in consumer_handles {
|
||||||
|
if let Err(e) = handle.await {
|
||||||
|
error!("消费者 task 结束异常: {:?}", e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
info!("队列系统已完全关闭");
|
||||||
|
}
|
||||||
@@ -0,0 +1,364 @@
|
|||||||
|
//! JSON 工具定义 — 像 curl + jq 一样描述一个工具
|
||||||
|
//!
|
||||||
|
//! 三要素:URL → 取节点 → 选字段
|
||||||
|
//!
|
||||||
|
//! ```json
|
||||||
|
//! {
|
||||||
|
//! "name": "stocks",
|
||||||
|
//! "description": "股票行情",
|
||||||
|
//! "url": "https://api.example.com/stocks",
|
||||||
|
//! "extract": "data.items",
|
||||||
|
//! "fields": ["name", "price"]
|
||||||
|
//! }
|
||||||
|
//! ```
|
||||||
|
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
|
/// HTTP 方法
|
||||||
|
#[derive(Debug, Clone, PartialEq)]
|
||||||
|
pub enum HttpMethod {
|
||||||
|
Get,
|
||||||
|
Post,
|
||||||
|
Put,
|
||||||
|
Delete,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[allow(dead_code)]
|
||||||
|
impl HttpMethod {
|
||||||
|
pub fn as_str(&self) -> &'static str {
|
||||||
|
match self {
|
||||||
|
HttpMethod::Get => "GET",
|
||||||
|
HttpMethod::Post => "POST",
|
||||||
|
HttpMethod::Put => "PUT",
|
||||||
|
HttpMethod::Delete => "DELETE",
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for HttpMethod {
|
||||||
|
fn default() -> Self {
|
||||||
|
HttpMethod::Get
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<'de> serde::Deserialize<'de> for HttpMethod {
|
||||||
|
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
|
||||||
|
where
|
||||||
|
D: serde::Deserializer<'de>,
|
||||||
|
{
|
||||||
|
let s = String::deserialize(deserializer)?;
|
||||||
|
match s.to_uppercase().as_str() {
|
||||||
|
"GET" => Ok(HttpMethod::Get),
|
||||||
|
"POST" => Ok(HttpMethod::Post),
|
||||||
|
"PUT" => Ok(HttpMethod::Put),
|
||||||
|
"DELETE" => Ok(HttpMethod::Delete),
|
||||||
|
other => Err(serde::de::Error::custom(format!(
|
||||||
|
"不支持的 HTTP 方法: {} (支持: GET/POST/PUT/DELETE)",
|
||||||
|
other
|
||||||
|
))),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 一个 JSON 工具定义
|
||||||
|
#[derive(Debug, Clone, serde::Deserialize)]
|
||||||
|
pub struct ToolDef {
|
||||||
|
/// 工具名称,用作 call_capability 的 name
|
||||||
|
pub name: String,
|
||||||
|
/// 工具描述
|
||||||
|
pub description: String,
|
||||||
|
/// URL 模板,如 "https://api.example.com/stock/{{code}}"
|
||||||
|
/// 支持 {{var}} 和 ${param:name:type} 两种传参形式
|
||||||
|
pub url: String,
|
||||||
|
/// HTTP 方法,默认 GET
|
||||||
|
#[serde(default)]
|
||||||
|
pub method: HttpMethod,
|
||||||
|
/// HTTP 头。值可包含 ${env:KEY}、${cmd:command}、${param:name:type}
|
||||||
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
|
pub headers: Option<HashMap<String, String>>,
|
||||||
|
/// POST/PUT 的 JSON body 模板。字符串值可包含 {{var}} 和 ${param:name:type}
|
||||||
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
|
pub body: Option<serde_json::Value>,
|
||||||
|
/// 显式参数声明(可选)
|
||||||
|
/// 格式:{ "name": { "type": "string", "description": "..." } }
|
||||||
|
/// 未声明的参数自动从 ${param:name:type} 和 {{var}} 推导
|
||||||
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
|
pub parameters: Option<serde_json::Value>,
|
||||||
|
/// JSON 提取路径(点路径),如 "data.items"
|
||||||
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
|
pub extract: Option<String>,
|
||||||
|
/// 筛选字段
|
||||||
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
|
pub fields: Option<Vec<String>>,
|
||||||
|
/// 超时秒数,默认 30
|
||||||
|
#[serde(default = "default_timeout")]
|
||||||
|
pub timeout: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
const fn default_timeout() -> u64 {
|
||||||
|
30
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ToolDef {
|
||||||
|
/// 从文件加载
|
||||||
|
pub fn from_file(path: &std::path::Path) -> Result<Self, String> {
|
||||||
|
let text = std::fs::read_to_string(path)
|
||||||
|
.map_err(|e| format!("读取 {} 失败: {}", path.display(), e))?;
|
||||||
|
serde_json::from_str(&text)
|
||||||
|
.map_err(|e| format!("解析 {} 失败: {}", path.display(), e))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 生成 SkillSpec
|
||||||
|
///
|
||||||
|
/// 参数推导优先级:
|
||||||
|
/// 1. 显式 parameters 字段(精确控制类型和描述)
|
||||||
|
/// 2. 从 URL/Body/Headers 中扫描 `${param:name:type}` 推导
|
||||||
|
/// 3. 从 {{var}} 推导(全视为 string 类型,兼容旧格式)
|
||||||
|
pub fn to_spec(&self) -> super::types::SkillSpec {
|
||||||
|
let params = self.infer_parameters();
|
||||||
|
super::types::SkillSpec {
|
||||||
|
name: self.name.clone(),
|
||||||
|
description: self.description.clone(),
|
||||||
|
risk_level: super::types::RiskLevel::Low,
|
||||||
|
parameters: params,
|
||||||
|
timeout_secs: self.timeout,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 推导参数定义
|
||||||
|
fn infer_parameters(&self) -> serde_json::Value {
|
||||||
|
// 1. 显式 parameters 优先
|
||||||
|
if let Some(ref p) = self.parameters {
|
||||||
|
let p = p.clone();
|
||||||
|
// 如果已经是完整 schema(含 type/properties),直接返回
|
||||||
|
if p.get("type").is_some() {
|
||||||
|
return p;
|
||||||
|
}
|
||||||
|
// 否则认为是属性定义 { name: { type, description }, ... }
|
||||||
|
// 包装为完整 schema
|
||||||
|
let required: Vec<String> = p.as_object()
|
||||||
|
.map(|o| o.keys().cloned().collect())
|
||||||
|
.unwrap_or_default();
|
||||||
|
return serde_json::json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": p,
|
||||||
|
"required": required
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// 2. 从模板中扫描 ${param:name:type} 和 {{var}}
|
||||||
|
let mut param_map: serde_json::Map<String, serde_json::Value> = serde_json::Map::new();
|
||||||
|
|
||||||
|
Self::collect_params(&self.url, &mut param_map);
|
||||||
|
if let Some(ref body) = self.body {
|
||||||
|
Self::collect_params_from_json(body, &mut param_map);
|
||||||
|
}
|
||||||
|
if let Some(ref headers) = self.headers {
|
||||||
|
for v in headers.values() {
|
||||||
|
Self::collect_params(v, &mut param_map);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let required: Vec<String> = param_map.keys().cloned().collect();
|
||||||
|
|
||||||
|
serde_json::json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": param_map,
|
||||||
|
"required": required
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 从字符串中收集参数(识别 ${param:name:type} 和 {{var}})
|
||||||
|
fn collect_params(s: &str, map: &mut serde_json::Map<String, serde_json::Value>) {
|
||||||
|
// 匹配 ${param:name:type} 和 ${param:name:type=default}
|
||||||
|
let mut pos = 0;
|
||||||
|
while let Some(start) = s[pos..].find("${param:") {
|
||||||
|
let rel = pos + start + 8; // skip "${param:"
|
||||||
|
let rest = &s[rel..];
|
||||||
|
if let Some(end) = rest.find('}') {
|
||||||
|
let segment = &rest[..end];
|
||||||
|
let parts: Vec<&str> = segment.split(':').collect();
|
||||||
|
let name = parts[0].trim().to_string();
|
||||||
|
if name.is_empty() {
|
||||||
|
pos = rel + end + 1;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if !map.contains_key(&name) {
|
||||||
|
let param_type = parts.get(1).copied().unwrap_or("string");
|
||||||
|
let description = parts.get(2).copied().unwrap_or(name.as_str());
|
||||||
|
let entry = serde_json::json!({
|
||||||
|
"type": param_type,
|
||||||
|
"description": description
|
||||||
|
});
|
||||||
|
map.insert(name.clone(), entry);
|
||||||
|
}
|
||||||
|
pos = rel + end + 1;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
pos = rel;
|
||||||
|
}
|
||||||
|
|
||||||
|
// 兼容旧格式 {{var}}
|
||||||
|
Self::extract_vars(s, map);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 从 JSON 值中递归收集参数
|
||||||
|
fn collect_params_from_json(
|
||||||
|
v: &serde_json::Value,
|
||||||
|
map: &mut serde_json::Map<String, serde_json::Value>,
|
||||||
|
) {
|
||||||
|
match v {
|
||||||
|
serde_json::Value::String(s) => Self::collect_params(s, map),
|
||||||
|
serde_json::Value::Object(obj) => {
|
||||||
|
for val in obj.values() {
|
||||||
|
Self::collect_params_from_json(val, map);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
serde_json::Value::Array(arr) => {
|
||||||
|
for item in arr {
|
||||||
|
Self::collect_params_from_json(item, map);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 兼容旧格式 {{var}} — 全视为 string 类型
|
||||||
|
fn extract_vars(s: &str, map: &mut serde_json::Map<String, serde_json::Value>) {
|
||||||
|
let mut remaining = s;
|
||||||
|
while let Some(start) = remaining.find("{{") {
|
||||||
|
remaining = &remaining[start + 2..];
|
||||||
|
if let Some(end) = remaining.find("}}") {
|
||||||
|
let name = remaining[..end].trim().to_string();
|
||||||
|
if !name.is_empty() && !name.starts_with('#') && !map.contains_key(&name) {
|
||||||
|
map.insert(
|
||||||
|
name.clone(),
|
||||||
|
serde_json::json!({"type": "string", "description": name}),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
remaining = &remaining[end + 2..];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_parse_minimal() {
|
||||||
|
let json = r#"{
|
||||||
|
"name": "ping",
|
||||||
|
"description": "Ping test",
|
||||||
|
"url": "https://httpbin.org/get"
|
||||||
|
}"#;
|
||||||
|
let def: ToolDef = serde_json::from_str(json).unwrap();
|
||||||
|
assert_eq!(def.name, "ping");
|
||||||
|
assert_eq!(def.method, HttpMethod::Get);
|
||||||
|
assert_eq!(def.timeout, 30);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_parse_post() {
|
||||||
|
let json = r#"{
|
||||||
|
"name": "translate",
|
||||||
|
"description": "翻译",
|
||||||
|
"url": "https://api.example.com/translate?text={{content}}",
|
||||||
|
"method": "POST",
|
||||||
|
"headers": { "Authorization": "${env:TR_KEY}" },
|
||||||
|
"body": { "q": "{{content}}", "target": "zh" },
|
||||||
|
"extract": "data.result",
|
||||||
|
"fields": ["text"],
|
||||||
|
"timeout": 60
|
||||||
|
}"#;
|
||||||
|
let def: ToolDef = serde_json::from_str(json).unwrap();
|
||||||
|
assert_eq!(def.method, HttpMethod::Post);
|
||||||
|
assert_eq!(def.timeout, 60);
|
||||||
|
assert_eq!(def.extract.as_deref(), Some("data.result"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_infer_parameters_from_url() {
|
||||||
|
let json = r#"{
|
||||||
|
"name": "test",
|
||||||
|
"description": "test",
|
||||||
|
"url": "https://api.example.com/{{city}}/weather?days={{days}}"
|
||||||
|
}"#;
|
||||||
|
let def: ToolDef = serde_json::from_str(json).unwrap();
|
||||||
|
let params = def.infer_parameters();
|
||||||
|
let props = params["properties"].as_object().unwrap();
|
||||||
|
assert!(props.contains_key("city"));
|
||||||
|
assert!(props.contains_key("days"));
|
||||||
|
let required = params["required"].as_array().unwrap();
|
||||||
|
assert!(required.iter().any(|v| v == "city"));
|
||||||
|
assert!(required.iter().any(|v| v == "days"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_infer_parameters_from_body() {
|
||||||
|
let json = r#"{
|
||||||
|
"name": "test",
|
||||||
|
"description": "test",
|
||||||
|
"url": "https://api.example.com/x",
|
||||||
|
"body": { "q": "{{text}}", "target": "{{lang}}" }
|
||||||
|
}"#;
|
||||||
|
let def: ToolDef = serde_json::from_str(json).unwrap();
|
||||||
|
let params = def.infer_parameters();
|
||||||
|
let props = params["properties"].as_object().unwrap();
|
||||||
|
assert!(props.contains_key("text"));
|
||||||
|
assert!(props.contains_key("lang"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_method_uppercase_flexibility() {
|
||||||
|
let json = r#"{"name":"t","description":"t","url":"u","method":"post"}"#;
|
||||||
|
let def: ToolDef = serde_json::from_str(json).unwrap();
|
||||||
|
assert_eq!(def.method, HttpMethod::Post);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_infer_parameters_from_param_syntax() {
|
||||||
|
let json = r#"{
|
||||||
|
"name": "weather",
|
||||||
|
"description": "天气查询",
|
||||||
|
"url": "https://api.example.com/weather/${param:city:string}",
|
||||||
|
"headers": { "X-Key": "${param:api_key:string}" }
|
||||||
|
}"#;
|
||||||
|
let def: ToolDef = serde_json::from_str(json).unwrap();
|
||||||
|
let params = def.infer_parameters();
|
||||||
|
let props = params["properties"].as_object().unwrap();
|
||||||
|
|
||||||
|
assert!(props.contains_key("city"), "应有 city 参数");
|
||||||
|
assert_eq!(props["city"]["type"], "string", "city 类型应为 string");
|
||||||
|
|
||||||
|
assert!(props.contains_key("api_key"), "应有 api_key 参数");
|
||||||
|
assert_eq!(props["api_key"]["type"], "string", "api_key 类型应为 string");
|
||||||
|
|
||||||
|
let required = params["required"].as_array().unwrap();
|
||||||
|
assert!(required.iter().any(|v| v == "city"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_infer_parameters_explicit_priority() {
|
||||||
|
let json = r#"{
|
||||||
|
"name": "my_tool",
|
||||||
|
"description": "test",
|
||||||
|
"url": "https://api.example.com/${param:city:string}",
|
||||||
|
"parameters": {
|
||||||
|
"city": { "type": "string", "description": "城市名称" },
|
||||||
|
"days": { "type": "integer", "description": "天数" }
|
||||||
|
}
|
||||||
|
}"#;
|
||||||
|
let def: ToolDef = serde_json::from_str(json).unwrap();
|
||||||
|
let params = def.infer_parameters();
|
||||||
|
|
||||||
|
// 显式 parameters 覆盖了自动推导
|
||||||
|
let props = params["properties"].as_object().unwrap();
|
||||||
|
assert!(props.contains_key("days"), "显式声明的 days 应在");
|
||||||
|
assert_eq!(props["days"]["type"], "integer");
|
||||||
|
|
||||||
|
// 且自动 wrap 为完整 schema
|
||||||
|
assert_eq!(params["type"], "object");
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,934 @@
|
|||||||
|
//! JSON 工具子进程执行器
|
||||||
|
//!
|
||||||
|
//! 从 ~/.ias/tools/*.json 加载 ToolDef,运行时通过 HTTP 执行。
|
||||||
|
//! 无需编译,放 JSON 即用。
|
||||||
|
//!
|
||||||
|
//! 执行流程:
|
||||||
|
//! 模板替换 → HTTP 请求 → JSON 解析
|
||||||
|
//! ├─ 失败 → 返回纯文本
|
||||||
|
//! ├─ 提取 → json_select()
|
||||||
|
//! ├─ 筛选 → filter_fields()
|
||||||
|
//! └─ 格式化输出
|
||||||
|
|
||||||
|
use crate::tools::definitions::{HttpMethod, ToolDef};
|
||||||
|
use crate::tools::types::{SkillResult, SkillSpec};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::path::{Path, PathBuf};
|
||||||
|
use std::sync::LazyLock;
|
||||||
|
use std::sync::Mutex;
|
||||||
|
use std::sync::RwLock;
|
||||||
|
use std::time::Instant;
|
||||||
|
|
||||||
|
/// 全局单例子进程执行器
|
||||||
|
static RUNNER: LazyLock<SubprocessRunner> = LazyLock::new(|| {
|
||||||
|
let dir = tools_dir();
|
||||||
|
SubprocessRunner::load_all(&dir)
|
||||||
|
});
|
||||||
|
|
||||||
|
/// 工具目录路径
|
||||||
|
fn tools_dir() -> std::path::PathBuf {
|
||||||
|
// 优先环境变量
|
||||||
|
if let Ok(dir) = std::env::var("IAS_TOOLS_DIR") {
|
||||||
|
return PathBuf::from(dir);
|
||||||
|
}
|
||||||
|
// 默认 ~/.ias/tools/
|
||||||
|
dirs::home_dir()
|
||||||
|
.unwrap_or_else(|| PathBuf::from("."))
|
||||||
|
.join(".ias")
|
||||||
|
.join("tools")
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 命令结果缓存
|
||||||
|
///
|
||||||
|
/// 避免 `${cmd:...}` 每次执行都重跑 shell 命令。
|
||||||
|
/// 默认 TTL 300 秒,适合 JWT token 等有时效的值。
|
||||||
|
struct CmdCache {
|
||||||
|
entries: Mutex<HashMap<String, (String, Instant)>>,
|
||||||
|
ttl_secs: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CmdCache {
|
||||||
|
fn new(ttl_secs: u64) -> Self {
|
||||||
|
Self {
|
||||||
|
entries: Mutex::new(HashMap::new()),
|
||||||
|
ttl_secs,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn get_or_exec(&self, cmd: &str) -> Result<String, String> {
|
||||||
|
// 检查缓存
|
||||||
|
{
|
||||||
|
let cache = self.entries.lock().unwrap();
|
||||||
|
if let Some((val, time)) = cache.get(cmd) {
|
||||||
|
if time.elapsed().as_secs() < self.ttl_secs {
|
||||||
|
return Ok(val.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 执行命令
|
||||||
|
let output = std::process::Command::new("bash")
|
||||||
|
.arg("-c")
|
||||||
|
.arg(cmd)
|
||||||
|
.stdout(std::process::Stdio::piped())
|
||||||
|
.stderr(std::process::Stdio::piped())
|
||||||
|
.output()
|
||||||
|
.map_err(|e| format!("执行命令失败: {}", e))?;
|
||||||
|
|
||||||
|
if !output.status.success() {
|
||||||
|
let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
|
||||||
|
return Err(format!("命令退出码 {}: {}", output.status.code().unwrap_or(-1), stderr));
|
||||||
|
}
|
||||||
|
|
||||||
|
let value = String::from_utf8_lossy(&output.stdout).trim().to_string();
|
||||||
|
if value.is_empty() {
|
||||||
|
return Err(format!("命令输出为空: {}", cmd));
|
||||||
|
}
|
||||||
|
|
||||||
|
// 写入缓存
|
||||||
|
{
|
||||||
|
let mut cache = self.entries.lock().unwrap();
|
||||||
|
cache.insert(cmd.to_string(), (value.clone(), Instant::now()));
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(value)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
static CMD_CACHE: LazyLock<CmdCache> = LazyLock::new(|| CmdCache::new(300));
|
||||||
|
|
||||||
|
/// 获取全局 runner 引用
|
||||||
|
pub fn runner() -> &'static SubprocessRunner {
|
||||||
|
&RUNNER
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 子进程执行器
|
||||||
|
pub struct SubprocessRunner {
|
||||||
|
tools: RwLock<HashMap<String, ToolDef>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl SubprocessRunner {
|
||||||
|
/// 加载目录下所有 *.json 文件
|
||||||
|
pub fn load_all(dir: &Path) -> Self {
|
||||||
|
let mut tools = HashMap::new();
|
||||||
|
if dir.exists() {
|
||||||
|
if let Ok(entries) = std::fs::read_dir(dir) {
|
||||||
|
for entry in entries.flatten() {
|
||||||
|
let path = entry.path();
|
||||||
|
if path.extension().map_or(false, |e| e == "json") {
|
||||||
|
match ToolDef::from_file(&path) {
|
||||||
|
Ok(def) => {
|
||||||
|
tracing::info!("加载工具: {} ← {}", def.name, path.display());
|
||||||
|
tools.insert(def.name.clone(), def);
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::warn!("跳过 {}: {}", path.display(), e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
tracing::info!("工具目录不存在,创建: {}", dir.display());
|
||||||
|
let _ = std::fs::create_dir_all(dir);
|
||||||
|
}
|
||||||
|
Self {
|
||||||
|
tools: RwLock::new(tools),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 重载全部工具(运行时热加载)
|
||||||
|
#[allow(dead_code)]
|
||||||
|
pub fn reload(&self) {
|
||||||
|
let dir = tools_dir();
|
||||||
|
let mut new_tools = HashMap::new();
|
||||||
|
if dir.exists() {
|
||||||
|
if let Ok(entries) = std::fs::read_dir(&dir) {
|
||||||
|
for entry in entries.flatten() {
|
||||||
|
let path = entry.path();
|
||||||
|
if path.extension().map_or(false, |e| e == "json") {
|
||||||
|
if let Ok(def) = ToolDef::from_file(&path) {
|
||||||
|
new_tools.insert(def.name.clone(), def);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if let Ok(mut w) = self.tools.write() {
|
||||||
|
*w = new_tools;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 是否有该工具
|
||||||
|
pub fn has_tool(&self, name: &str) -> bool {
|
||||||
|
self.tools.read().map_or(false, |m| m.contains_key(name))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 获取工具 spec(参数从 URL/Body 模板推导)
|
||||||
|
#[allow(dead_code)]
|
||||||
|
pub fn spec(&self, name: &str) -> Option<SkillSpec> {
|
||||||
|
self.tools.read().ok().and_then(|m| m.get(name).map(|d| d.to_spec()))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 获取所有工具 specs
|
||||||
|
pub fn all_specs(&self) -> Vec<SkillSpec> {
|
||||||
|
self.tools
|
||||||
|
.read()
|
||||||
|
.map(|m| m.values().map(|d| d.to_spec()).collect())
|
||||||
|
.unwrap_or_default()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 执行工具
|
||||||
|
pub async fn execute(&self, name: &str, args_json: &str) -> Option<SkillResult> {
|
||||||
|
let def = self.tools.read().ok()?.get(name)?.clone();
|
||||||
|
let args: serde_json::Value = serde_json::from_str(args_json).unwrap_or_default();
|
||||||
|
Some(self.execute_def(&def, &args).await)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 执行一个 ToolDef
|
||||||
|
async fn execute_def(&self, def: &ToolDef, args: &serde_json::Value) -> SkillResult {
|
||||||
|
// 1. 模板替换
|
||||||
|
let url = render_url(&def.url, args);
|
||||||
|
let headers = resolve_headers_with_args(def.headers.as_ref(), args);
|
||||||
|
let body = def.body.as_ref().map(|b| render_body(b, args));
|
||||||
|
|
||||||
|
// 2. HTTP 请求
|
||||||
|
let resp_text = match http_request(def.method.clone(), &url, &headers, body.as_deref()).await
|
||||||
|
{
|
||||||
|
Ok(t) => t,
|
||||||
|
Err(e) => return SkillResult::error(format!("请求失败: {}", e)),
|
||||||
|
};
|
||||||
|
|
||||||
|
// 3. JSON 判定
|
||||||
|
let json: serde_json::Value = match serde_json::from_str(&resp_text) {
|
||||||
|
Ok(j) => j,
|
||||||
|
Err(_) => {
|
||||||
|
// 纯文本,直接返回全文
|
||||||
|
return SkillResult::ok(resp_text);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// 4. 提取子节点
|
||||||
|
let extracted = match &def.extract {
|
||||||
|
Some(path) => match json_select(&json, path) {
|
||||||
|
Some(v) => v.clone(),
|
||||||
|
None => {
|
||||||
|
return SkillResult::error(format!("路径 '{}' 未找到", path));
|
||||||
|
}
|
||||||
|
},
|
||||||
|
None => json,
|
||||||
|
};
|
||||||
|
|
||||||
|
// 5. 筛选字段
|
||||||
|
let filtered = match &def.fields {
|
||||||
|
Some(fields) => filter_fields(&extracted, fields),
|
||||||
|
None => extracted,
|
||||||
|
};
|
||||||
|
|
||||||
|
// 6. 格式化输出
|
||||||
|
let output = format_output(&filtered);
|
||||||
|
|
||||||
|
SkillResult::ok_with_data(output, filtered)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ═══════════════════════════════════════════════
|
||||||
|
// 模板渲染
|
||||||
|
// ═══════════════════════════════════════════════
|
||||||
|
|
||||||
|
/// 渲染 URL 模板:{{var}} / ${param:name:type} → args.name(URL 编码)
|
||||||
|
fn render_url(template: &str, args: &serde_json::Value) -> String {
|
||||||
|
// 先处理 ${param:name:type}
|
||||||
|
let after_param = replace_param_refs(template, args, true);
|
||||||
|
// 再处理 {{var}}
|
||||||
|
replace_mustache(&after_param, args)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 渲染普通文本(非 URL 编码),支持三种引用:
|
||||||
|
/// ${param:name:type} — 参数值
|
||||||
|
/// ${env:KEY} — 环境变量
|
||||||
|
/// ${cmd:command} — shell 命令输出
|
||||||
|
fn render_text(template: &str, args: &serde_json::Value) -> String {
|
||||||
|
// 先处理 ${param:name:type}
|
||||||
|
let after_param = replace_param_refs(template, args, false);
|
||||||
|
// 再处理 {{var}}(不编码)
|
||||||
|
replace_mustache_plain(&after_param, args)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 替换 ${param:name:type} 引用
|
||||||
|
fn replace_param_refs(template: &str, args: &serde_json::Value, url_encode: bool) -> String {
|
||||||
|
let mut result = String::new();
|
||||||
|
let mut remaining = template;
|
||||||
|
|
||||||
|
while let Some(start) = remaining.find("${param:") {
|
||||||
|
result.push_str(&remaining[..start]);
|
||||||
|
let rest = &remaining[start + 8..];
|
||||||
|
if let Some(end) = rest.find('}') {
|
||||||
|
let segment = &rest[..end];
|
||||||
|
let name = segment.split(':').next().map(|s| s.trim()).unwrap_or("");
|
||||||
|
if !name.is_empty() {
|
||||||
|
let value = args
|
||||||
|
.get(name)
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.unwrap_or("");
|
||||||
|
if url_encode {
|
||||||
|
result.push_str(&urlencoding(value));
|
||||||
|
} else {
|
||||||
|
result.push_str(value);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
remaining = &rest[end + 1..];
|
||||||
|
} else {
|
||||||
|
result.push_str(&remaining[start..]);
|
||||||
|
remaining = "";
|
||||||
|
}
|
||||||
|
}
|
||||||
|
result.push_str(remaining);
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 替换 {{var}} — URL 编码版
|
||||||
|
fn replace_mustache(template: &str, args: &serde_json::Value) -> String {
|
||||||
|
let mut result = String::new();
|
||||||
|
let mut remaining = template;
|
||||||
|
while let Some(start) = remaining.find("{{") {
|
||||||
|
result.push_str(&remaining[..start]);
|
||||||
|
remaining = &remaining[start + 2..];
|
||||||
|
if let Some(end) = remaining.find("}}") {
|
||||||
|
let key = remaining[..end].trim();
|
||||||
|
if !key.starts_with('#') {
|
||||||
|
let value = args
|
||||||
|
.get(key)
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.map(|s| urlencoding(s))
|
||||||
|
.unwrap_or_else(|| key.to_string());
|
||||||
|
result.push_str(&value);
|
||||||
|
}
|
||||||
|
remaining = &remaining[end + 2..];
|
||||||
|
} else {
|
||||||
|
result.push_str(&remaining[..]);
|
||||||
|
remaining = "";
|
||||||
|
}
|
||||||
|
}
|
||||||
|
result.push_str(remaining);
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 替换 {{var}} — 纯文本版(不编码)
|
||||||
|
fn replace_mustache_plain(template: &str, args: &serde_json::Value) -> String {
|
||||||
|
let mut result = String::new();
|
||||||
|
let mut remaining = template;
|
||||||
|
while let Some(start) = remaining.find("{{") {
|
||||||
|
result.push_str(&remaining[..start]);
|
||||||
|
remaining = &remaining[start + 2..];
|
||||||
|
if let Some(end) = remaining.find("}}") {
|
||||||
|
let key = remaining[..end].trim();
|
||||||
|
if !key.starts_with('#') {
|
||||||
|
let value = args
|
||||||
|
.get(key)
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.unwrap_or(key);
|
||||||
|
result.push_str(value);
|
||||||
|
}
|
||||||
|
remaining = &remaining[end + 2..];
|
||||||
|
} else {
|
||||||
|
result.push_str(&remaining[..]);
|
||||||
|
remaining = "";
|
||||||
|
}
|
||||||
|
}
|
||||||
|
result.push_str(remaining);
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
/// URL 编码
|
||||||
|
fn urlencoding(s: &str) -> String {
|
||||||
|
// 只编码需要编码的字符
|
||||||
|
let mut result = String::with_capacity(s.len());
|
||||||
|
for c in s.chars() {
|
||||||
|
match c {
|
||||||
|
'A'..='Z' | 'a'..='z' | '0'..='9' | '-' | '_' | '.' | '~' => result.push(c),
|
||||||
|
_ => {
|
||||||
|
for byte in c.to_string().as_bytes() {
|
||||||
|
result.push_str(&format!("%{:02X}", byte));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 解析模板值(无 args 上下文,主要用于测试)
|
||||||
|
#[allow(dead_code)]
|
||||||
|
fn resolve_template_value(template: &str) -> String {
|
||||||
|
if let Some(env_key) = template.strip_prefix("${env:").and_then(|s| s.strip_suffix('}')) {
|
||||||
|
return std::env::var(env_key).unwrap_or_default();
|
||||||
|
}
|
||||||
|
if let Some(cmd) = template.strip_prefix("${cmd:").and_then(|s| s.strip_suffix('}')) {
|
||||||
|
return CMD_CACHE.get_or_exec(cmd).unwrap_or_else(|e| {
|
||||||
|
tracing::warn!("执行命令失败: {}", e);
|
||||||
|
String::new()
|
||||||
|
});
|
||||||
|
}
|
||||||
|
// ${param:name:type} — 在 resolve_headers 中通过 args 解析,
|
||||||
|
// 单纯 resolve_template_value 没有 args 上下文的,先保留原样
|
||||||
|
if template.contains("${param:") {
|
||||||
|
return template.to_string();
|
||||||
|
}
|
||||||
|
template.to_string()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 渲染 header 值(带 args 上下文,支持 ${param:name:type})
|
||||||
|
fn resolve_header_value(template: &str, args: &serde_json::Value) -> String {
|
||||||
|
// 优先 env/cmd
|
||||||
|
if let Some(env_key) = template.strip_prefix("${env:").and_then(|s| s.strip_suffix('}')) {
|
||||||
|
return std::env::var(env_key).unwrap_or_default();
|
||||||
|
}
|
||||||
|
if let Some(cmd) = template.strip_prefix("${cmd:").and_then(|s| s.strip_suffix('}')) {
|
||||||
|
return CMD_CACHE.get_or_exec(cmd).unwrap_or_else(|e| {
|
||||||
|
tracing::warn!("执行命令失败: {}", e);
|
||||||
|
String::new()
|
||||||
|
});
|
||||||
|
}
|
||||||
|
// ${param:name:type} → 取 args.name
|
||||||
|
if let Some(param_rest) = template.strip_prefix("${param:") {
|
||||||
|
if let Some(end) = param_rest.find('}') {
|
||||||
|
let segment = ¶m_rest[..end];
|
||||||
|
let name = segment.split(':').next().map(|s| s.trim()).unwrap_or("");
|
||||||
|
if !name.is_empty() {
|
||||||
|
return args
|
||||||
|
.get(name)
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.unwrap_or("")
|
||||||
|
.to_string();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
render_text(template, args)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 渲染 header 列表(带 args)
|
||||||
|
fn resolve_headers_with_args(
|
||||||
|
headers: Option<&HashMap<String, String>>,
|
||||||
|
args: &serde_json::Value,
|
||||||
|
) -> Vec<(String, String)> {
|
||||||
|
headers
|
||||||
|
.map(|h| {
|
||||||
|
h.iter()
|
||||||
|
.map(|(k, v)| (k.clone(), resolve_header_value(v, args)))
|
||||||
|
.collect()
|
||||||
|
})
|
||||||
|
.unwrap_or_default()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 渲染 header 值(无 args 上下文,主要用于测试)
|
||||||
|
#[allow(dead_code)]
|
||||||
|
fn resolve_headers(
|
||||||
|
headers: Option<&HashMap<String, String>>,
|
||||||
|
) -> Vec<(String, String)> {
|
||||||
|
headers
|
||||||
|
.map(|h| {
|
||||||
|
h.iter()
|
||||||
|
.map(|(k, v)| (k.clone(), resolve_template_value(v)))
|
||||||
|
.collect()
|
||||||
|
})
|
||||||
|
.unwrap_or_default()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 渲染 JSON body 模板:字符串中的 {{var}} 和 ${param:name:type} 被替换
|
||||||
|
fn render_body(template: &serde_json::Value, args: &serde_json::Value) -> String {
|
||||||
|
fn render_value(v: &serde_json::Value, args: &serde_json::Value) -> serde_json::Value {
|
||||||
|
match v {
|
||||||
|
serde_json::Value::String(s) => {
|
||||||
|
if s.contains("{{") || s.contains("${param:") {
|
||||||
|
let rendered = render_text(s, args);
|
||||||
|
// 尝试解析为 JSON
|
||||||
|
if let Ok(j) = serde_json::from_str::<serde_json::Value>(&rendered) {
|
||||||
|
return j;
|
||||||
|
}
|
||||||
|
return serde_json::Value::String(rendered);
|
||||||
|
}
|
||||||
|
v.clone()
|
||||||
|
}
|
||||||
|
serde_json::Value::Object(obj) => {
|
||||||
|
let mut m = serde_json::Map::new();
|
||||||
|
for (k, val) in obj {
|
||||||
|
m.insert(k.clone(), render_value(val, args));
|
||||||
|
}
|
||||||
|
serde_json::Value::Object(m)
|
||||||
|
}
|
||||||
|
serde_json::Value::Array(arr) => {
|
||||||
|
let items: Vec<_> = arr.iter().map(|item| render_value(item, args)).collect();
|
||||||
|
serde_json::Value::Array(items)
|
||||||
|
}
|
||||||
|
_ => v.clone(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let rendered = render_value(template, args);
|
||||||
|
serde_json::to_string(&rendered).unwrap_or_default()
|
||||||
|
}
|
||||||
|
|
||||||
|
// ═══════════════════════════════════════════════
|
||||||
|
// HTTP 请求
|
||||||
|
// ═══════════════════════════════════════════════
|
||||||
|
|
||||||
|
async fn http_request(
|
||||||
|
method: HttpMethod,
|
||||||
|
url: &str,
|
||||||
|
headers: &[(String, String)],
|
||||||
|
body: Option<&str>,
|
||||||
|
) -> Result<String, String> {
|
||||||
|
let client = reqwest::Client::builder()
|
||||||
|
.timeout(std::time::Duration::from_secs(30))
|
||||||
|
.user_agent("iAs/1.0 (JSON Tool Runner)")
|
||||||
|
.build()
|
||||||
|
.map_err(|e| format!("创建 HTTP client 失败: {}", e))?;
|
||||||
|
|
||||||
|
let mut req = match method {
|
||||||
|
HttpMethod::Get => client.get(url),
|
||||||
|
HttpMethod::Post => {
|
||||||
|
let mut r = client.post(url);
|
||||||
|
if let Some(b) = body {
|
||||||
|
r = r.header("Content-Type", "application/json").body(b.to_string());
|
||||||
|
}
|
||||||
|
r
|
||||||
|
}
|
||||||
|
HttpMethod::Put => {
|
||||||
|
let mut r = client.put(url);
|
||||||
|
if let Some(b) = body {
|
||||||
|
r = r.header("Content-Type", "application/json").body(b.to_string());
|
||||||
|
}
|
||||||
|
r
|
||||||
|
}
|
||||||
|
HttpMethod::Delete => client.delete(url),
|
||||||
|
};
|
||||||
|
|
||||||
|
for (k, v) in headers {
|
||||||
|
if !v.is_empty() {
|
||||||
|
req = req.header(k.as_str(), v.as_str());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let resp = req.send().await.map_err(|e| format!("HTTP 请求失败: {}", e))?;
|
||||||
|
|
||||||
|
if !resp.status().is_success() {
|
||||||
|
let status = resp.status();
|
||||||
|
let text = resp.text().await.unwrap_or_default();
|
||||||
|
return Err(format!("HTTP {}: {}", status, text.chars().take(200).collect::<String>()));
|
||||||
|
}
|
||||||
|
|
||||||
|
resp.text().await.map_err(|e| format!("读取响应失败: {}", e))
|
||||||
|
}
|
||||||
|
|
||||||
|
// ═══════════════════════════════════════════════
|
||||||
|
// JSON 节点选择(点路径)
|
||||||
|
// ═══════════════════════════════════════════════
|
||||||
|
|
||||||
|
/// 按点路径选择 JSON 节点
|
||||||
|
///
|
||||||
|
/// 支持格式:
|
||||||
|
/// "data" → json["data"]
|
||||||
|
/// "data.items" → json["data"]["items"]
|
||||||
|
/// "items.0.name" → json["items"][0]["name"]
|
||||||
|
/// ".data.items" → 自动忽略前导点
|
||||||
|
fn json_select<'a>(json: &'a serde_json::Value, path: &str) -> Option<&'a serde_json::Value> {
|
||||||
|
let path = path.strip_prefix('.').unwrap_or(path);
|
||||||
|
if path.is_empty() {
|
||||||
|
return Some(json);
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut current = json;
|
||||||
|
for segment in path.split('.') {
|
||||||
|
if segment.is_empty() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
current = match current {
|
||||||
|
serde_json::Value::Object(obj) => obj.get(segment)?,
|
||||||
|
serde_json::Value::Array(arr) => {
|
||||||
|
let idx: usize = segment.parse().ok()?;
|
||||||
|
arr.get(idx)?
|
||||||
|
}
|
||||||
|
_ => return None,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
Some(current)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ═══════════════════════════════════════════════
|
||||||
|
// 字段筛选
|
||||||
|
// ═══════════════════════════════════════════════
|
||||||
|
|
||||||
|
/// 只保留指定字段
|
||||||
|
///
|
||||||
|
/// - object: 只保留 fields 中的 key
|
||||||
|
/// - array: 对每个元素做同样筛选
|
||||||
|
/// - 其他: 保留原值
|
||||||
|
fn filter_fields(value: &serde_json::Value, fields: &[String]) -> serde_json::Value {
|
||||||
|
match value {
|
||||||
|
serde_json::Value::Object(obj) => {
|
||||||
|
let mut m = serde_json::Map::new();
|
||||||
|
for field in fields {
|
||||||
|
if let Some(v) = obj.get(field) {
|
||||||
|
m.insert(field.clone(), v.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
serde_json::Value::Object(m)
|
||||||
|
}
|
||||||
|
serde_json::Value::Array(arr) => {
|
||||||
|
let items: Vec<_> = arr.iter().map(|item| filter_fields(item, fields)).collect();
|
||||||
|
serde_json::Value::Array(items)
|
||||||
|
}
|
||||||
|
other => other.clone(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ═══════════════════════════════════════════════
|
||||||
|
// 输出格式化
|
||||||
|
// ═══════════════════════════════════════════════
|
||||||
|
|
||||||
|
/// 将提取/筛选后的结果格式化为可读文本
|
||||||
|
fn format_output(value: &serde_json::Value) -> String {
|
||||||
|
match value {
|
||||||
|
serde_json::Value::Null => "null".to_string(),
|
||||||
|
serde_json::Value::Bool(b) => b.to_string(),
|
||||||
|
serde_json::Value::Number(n) => n.to_string(),
|
||||||
|
serde_json::Value::String(s) => s.clone(),
|
||||||
|
serde_json::Value::Array(arr) => format_array(arr),
|
||||||
|
serde_json::Value::Object(obj) => format_object(obj),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn format_array(arr: &[serde_json::Value]) -> String {
|
||||||
|
if arr.is_empty() {
|
||||||
|
return "(空)".to_string();
|
||||||
|
}
|
||||||
|
|
||||||
|
let lines: Vec<String> = arr
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(i, item)| {
|
||||||
|
let prefix = format!("{}. ", i + 1);
|
||||||
|
let body = match item {
|
||||||
|
serde_json::Value::Object(_) => {
|
||||||
|
// 对象数组:每个对象展平成 key: value 行
|
||||||
|
let obj_lines = format_object_lines(item)
|
||||||
|
.into_iter()
|
||||||
|
.map(|l| format!(" {}", l))
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
.join("\n");
|
||||||
|
obj_lines
|
||||||
|
}
|
||||||
|
serde_json::Value::Array(a) => format_array(a),
|
||||||
|
other => format_output(other),
|
||||||
|
};
|
||||||
|
format!("{}{}", prefix, body)
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
lines.join("\n---\n")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn format_object(obj: &serde_json::Map<String, serde_json::Value>) -> String {
|
||||||
|
let lines = format_object_lines(&serde_json::Value::Object(obj.clone()));
|
||||||
|
lines.join("\n")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn format_object_lines(value: &serde_json::Value) -> Vec<String> {
|
||||||
|
match value {
|
||||||
|
serde_json::Value::Object(obj) => {
|
||||||
|
let mut lines = Vec::new();
|
||||||
|
for (k, v) in obj {
|
||||||
|
let v_str = match v {
|
||||||
|
serde_json::Value::String(s) => s.clone(),
|
||||||
|
serde_json::Value::Null => "null".to_string(),
|
||||||
|
serde_json::Value::Bool(b) => b.to_string(),
|
||||||
|
serde_json::Value::Number(n) => n.to_string(),
|
||||||
|
serde_json::Value::Array(arr) => {
|
||||||
|
if arr.len() <= 5 {
|
||||||
|
arr.iter()
|
||||||
|
.map(format_output)
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
.join(", ")
|
||||||
|
} else {
|
||||||
|
format!("[{} 项]", arr.len())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
serde_json::Value::Object(o) => {
|
||||||
|
format!("{{ {} }}", o.keys().cloned().collect::<Vec<_>>().join(", "))
|
||||||
|
}
|
||||||
|
};
|
||||||
|
lines.push(format!(" {}: {}", k, v_str));
|
||||||
|
}
|
||||||
|
lines
|
||||||
|
}
|
||||||
|
_ => vec![format_output(value)],
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
// ─── 模板渲染 ───
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_render_url_basic() {
|
||||||
|
let args = serde_json::json!({"code": "000001"});
|
||||||
|
let url = render_url("https://api.example.com/stock/{{code}}", &args);
|
||||||
|
assert_eq!(url, "https://api.example.com/stock/000001");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_render_url_encode_chinese() {
|
||||||
|
let args = serde_json::json!({"city": "北京"});
|
||||||
|
let url = render_url("https://api.example.com/weather?city={{city}}", &args);
|
||||||
|
assert!(url.contains("%E5%8C%97%E4%BA%AC"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_render_url_no_vars() {
|
||||||
|
let args = serde_json::json!({});
|
||||||
|
let url = render_url("https://api.example.com/ping", &args);
|
||||||
|
assert_eq!(url, "https://api.example.com/ping");
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── JSON 选择 ───
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_json_select_simple() {
|
||||||
|
let json = serde_json::json!({"data": {"price": 42.5}});
|
||||||
|
assert_eq!(json_select(&json, "data.price").and_then(|v| v.as_f64()), Some(42.5));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_json_select_dot_prefix() {
|
||||||
|
let json = serde_json::json!({"data": {"price": 42.5}});
|
||||||
|
assert_eq!(json_select(&json, ".data.price").and_then(|v| v.as_f64()), Some(42.5));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_json_select_array_index() {
|
||||||
|
let json = serde_json::json!({"items": [{"name": "a"}, {"name": "b"}]});
|
||||||
|
assert_eq!(
|
||||||
|
json_select(&json, "items.1.name").and_then(|v| v.as_str()),
|
||||||
|
Some("b")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_json_select_not_found() {
|
||||||
|
let json = serde_json::json!({"a": 1});
|
||||||
|
assert!(json_select(&json, "a.b.c").is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── 字段筛选 ───
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_filter_fields_object() {
|
||||||
|
let obj = serde_json::json!({"name": "AAPL", "price": 150, "ignored": true});
|
||||||
|
let fields = vec!["name".into(), "price".into()];
|
||||||
|
let filtered = filter_fields(&obj, &fields);
|
||||||
|
assert_eq!(filtered.as_object().unwrap().len(), 2);
|
||||||
|
assert!(filtered.get("ignored").is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_filter_fields_array() {
|
||||||
|
let arr = serde_json::json!([
|
||||||
|
{"name": "A", "price": 10, "extra": true},
|
||||||
|
{"name": "B", "price": 20, "extra": false}
|
||||||
|
]);
|
||||||
|
let fields = vec!["name".into()];
|
||||||
|
let filtered = filter_fields(&arr, &fields);
|
||||||
|
let items = filtered.as_array().unwrap();
|
||||||
|
assert_eq!(items.len(), 2);
|
||||||
|
assert!(items[0].get("price").is_none());
|
||||||
|
assert!(items[0].get("name").is_some());
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── 输出格式化 ───
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_format_scalar() {
|
||||||
|
assert_eq!(format_output(&serde_json::json!(42)), "42");
|
||||||
|
assert_eq!(format_output(&serde_json::json!("hello")), "hello");
|
||||||
|
assert_eq!(format_output(&serde_json::json!(true)), "true");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_format_object() {
|
||||||
|
let obj = serde_json::json!({"name": "AAPL", "price": 150.5});
|
||||||
|
let out = format_output(&obj);
|
||||||
|
assert!(out.contains("name: AAPL"));
|
||||||
|
assert!(out.contains("price: 150.5"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_format_array_of_objects() {
|
||||||
|
let arr = serde_json::json!([
|
||||||
|
{"name": "A", "price": 10},
|
||||||
|
{"name": "B", "price": 20}
|
||||||
|
]);
|
||||||
|
let out = format_output(&arr);
|
||||||
|
assert!(out.starts_with("1."));
|
||||||
|
assert!(out.contains("---"));
|
||||||
|
assert!(out.contains("2."));
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── 头部解析 ───
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_resolve_headers_static() {
|
||||||
|
let mut h = HashMap::new();
|
||||||
|
h.insert("Accept".into(), "application/json".into());
|
||||||
|
let resolved = resolve_headers(Some(&h));
|
||||||
|
assert_eq!(resolved[0].1, "application/json");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_resolve_cmd_inline() {
|
||||||
|
let result = resolve_template_value("${cmd:echo hello-world}");
|
||||||
|
assert_eq!(result, "hello-world");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_resolve_cmd_cached() {
|
||||||
|
// 同一命令应缓存,第二次直接返回
|
||||||
|
let r1 = resolve_template_value("${cmd:echo cached-test-123}");
|
||||||
|
let r2 = resolve_template_value("${cmd:echo cached-test-123}");
|
||||||
|
assert_eq!(r1, "cached-test-123");
|
||||||
|
assert_eq!(r2, r1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_resolve_env() {
|
||||||
|
unsafe { std::env::set_var("TEST_API_KEY", "secret123"); }
|
||||||
|
let mut h = HashMap::new();
|
||||||
|
h.insert("X-Key".into(), "${env:TEST_API_KEY}".into());
|
||||||
|
let resolved = resolve_headers(Some(&h));
|
||||||
|
assert_eq!(resolved[0].1, "secret123");
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── 新格式 ${param:name:type} 渲染 ───
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_render_url_with_param() {
|
||||||
|
let args = serde_json::json!({"city": "北京", "days": "3"});
|
||||||
|
let url = render_url(
|
||||||
|
"https://api.example.com/weather/${param:city:string}/days/${param:days:integer}",
|
||||||
|
&args,
|
||||||
|
);
|
||||||
|
assert!(url.contains("%E5%8C%97%E4%BA%AC"));
|
||||||
|
assert!(url.contains("3"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_render_url_mixed() {
|
||||||
|
let args = serde_json::json!({"code": "000001"});
|
||||||
|
let url = render_url(
|
||||||
|
"https://api.example.com/${param:code:string}?type={{#comment}}",
|
||||||
|
&args,
|
||||||
|
);
|
||||||
|
// {{#comment}} 被跳过(注释语法)
|
||||||
|
assert_eq!(url, "https://api.example.com/000001?type=");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_render_text_param() {
|
||||||
|
let args = serde_json::json!({"name": "hello"});
|
||||||
|
let result = render_text("prefix_${param:name:string}_suffix", &args);
|
||||||
|
assert_eq!(result, "prefix_hello_suffix");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_resolve_header_param() {
|
||||||
|
let args = serde_json::json!({"token": "abc123"});
|
||||||
|
let mut h = HashMap::new();
|
||||||
|
h.insert("Authorization".into(), "Bearer ${param:token:string}".into());
|
||||||
|
let resolved = resolve_headers_with_args(Some(&h), &args);
|
||||||
|
assert_eq!(resolved[0].1, "Bearer abc123");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_param_in_body() {
|
||||||
|
let args = serde_json::json!({"q": "你好", "lang": "zh"});
|
||||||
|
let body = serde_json::json!({
|
||||||
|
"query": "${param:q:string}",
|
||||||
|
"language": "${param:lang:string}"
|
||||||
|
});
|
||||||
|
let rendered = render_body(&body, &args);
|
||||||
|
assert!(rendered.contains("你好"));
|
||||||
|
assert!(rendered.contains("zh"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_load_tools_from_disk() {
|
||||||
|
let dir = std::path::Path::new("/home/xiao/.ias/tools");
|
||||||
|
if !dir.exists() {
|
||||||
|
eprintln!("跳过磁盘加载测试: 目录不存在");
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let runner = SubprocessRunner::load_all(dir);
|
||||||
|
assert!(runner.has_tool("weather"), "应加载 weather 工具");
|
||||||
|
assert!(runner.has_tool("github_repo"), "应加载 github_repo 工具");
|
||||||
|
|
||||||
|
let specs = runner.all_specs();
|
||||||
|
assert!(specs.len() >= 2, "至少加载 2 个工具, 实际: {}", specs.len());
|
||||||
|
|
||||||
|
// 验证 weather 工具的参数从模板推导
|
||||||
|
let weather_spec = runner.spec("weather").unwrap();
|
||||||
|
let params = &weather_spec.parameters;
|
||||||
|
let props = params["properties"].as_object().unwrap();
|
||||||
|
assert!(props.contains_key("city"), "weather 应有 city 参数");
|
||||||
|
}
|
||||||
|
|
||||||
|
// ─── 完整端到端 ───
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_execute_json_weather() {
|
||||||
|
let dir = std::path::Path::new("/home/xiao/.ias/tools");
|
||||||
|
if !dir.exists() {
|
||||||
|
eprintln!("跳过: 工具目录不存在");
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let runner = SubprocessRunner::load_all(dir);
|
||||||
|
let result = runner.execute("weather", r#"{"city":"北京"}"#).await.unwrap();
|
||||||
|
assert!(result.success, "天气查询应成功: {}", result.output);
|
||||||
|
println!("🌤️ 天气结果: {}", result.output);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_execute_json_github() {
|
||||||
|
let dir = std::path::Path::new("/home/xiao/.ias/tools");
|
||||||
|
if !dir.exists() {
|
||||||
|
eprintln!("跳过: 工具目录不存在");
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let runner = SubprocessRunner::load_all(dir);
|
||||||
|
let result = runner.execute("github_repo", r#"{"owner":"rust-lang","repo":"rust"}"#).await.unwrap();
|
||||||
|
assert!(result.success, "GitHub 查询应成功: {}", result.output);
|
||||||
|
assert!(result.output.contains("rust"), "应有仓库名: {}", result.output);
|
||||||
|
println!("🐙 GitHub: {}", result.output);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_execute_httpbin_get() {
|
||||||
|
let def = ToolDef {
|
||||||
|
name: "httpbin".into(),
|
||||||
|
description: "test".into(),
|
||||||
|
url: "https://httpbin.org/get".into(),
|
||||||
|
method: HttpMethod::Get,
|
||||||
|
headers: None,
|
||||||
|
body: None,
|
||||||
|
parameters: None,
|
||||||
|
extract: Some("url".into()),
|
||||||
|
fields: None,
|
||||||
|
timeout: 30,
|
||||||
|
};
|
||||||
|
let args = serde_json::json!({});
|
||||||
|
let runner = SubprocessRunner::load_all(Path::new("/nonexistent"));
|
||||||
|
let result = runner.execute_def(&def, &args).await;
|
||||||
|
assert!(result.success, "HTTP 请求应成功: {}", result.output);
|
||||||
|
assert!(result.output.contains("httpbin.org/get"), "应提取 url 字段: {}", result.output);
|
||||||
|
}
|
||||||
|
}
|
||||||
+7
-3
@@ -1,6 +1,9 @@
|
|||||||
//! Worker — 无状态执行进程
|
#![deprecated = "Worker 进程模式已废弃,请使用 daemon 模式(统一 tokio task 架构)"]
|
||||||
|
|
||||||
|
//! Worker — 无状态执行进程(已废弃)
|
||||||
//!
|
//!
|
||||||
//! 每条消息 spawn 一个 worker:
|
//! 此模块保留供编译参考,不再被 daemon 使用。
|
||||||
|
//! 新架构使用 daemon 内的 LLM/Tool/Send 三个消费者 tokio task。
|
||||||
//! 1. 连接 daemon 的 Unix Domain Socket
|
//! 1. 连接 daemon 的 Unix Domain Socket
|
||||||
//! 2. 读取 TaskFrame (用户消息 + 上下文 + 环境变量)
|
//! 2. 读取 TaskFrame (用户消息 + 上下文 + 环境变量)
|
||||||
//! 3. 注入环境变量,构建 Conversation,执行 LLM 对话 + 工具调用
|
//! 3. 注入环境变量,构建 Conversation,执行 LLM 对话 + 工具调用
|
||||||
@@ -103,7 +106,8 @@ pub async fn run(sock_path: &str) -> Result<(), String> {
|
|||||||
conv.set_tool_executor(executor);
|
conv.set_tool_executor(executor);
|
||||||
|
|
||||||
// 7. 执行 LLM 对话(带审批检测)
|
// 7. 执行 LLM 对话(带审批检测)
|
||||||
let (reply, _used_tools, usage) = conv
|
use crate::llm::ChatResult;
|
||||||
|
let ChatResult { reply, used_tools: _used_tools, usage, .. } = conv
|
||||||
.chat_with_tools(task.msg.text.clone())
|
.chat_with_tools(task.msg.text.clone())
|
||||||
.await
|
.await
|
||||||
.map_err(|e| format!("LLM 对话失败: {}", e))?;
|
.map_err(|e| format!("LLM 对话失败: {}", e))?;
|
||||||
|
|||||||
Reference in New Issue
Block a user