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:
2026-06-09 16:32:49 +08:00
parent 8dc232ebab
commit 937556917c
17 changed files with 2937 additions and 729 deletions
+10
View File
@@ -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"] }
+143
View File
@@ -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()),
}
}
+104
View File
@@ -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())
))
}
+38
View File
@@ -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
View File
@@ -115,7 +115,7 @@ pub enum Commands {
sock: String,
},
/// Worker 进程( daemon 自动 spawn,也可手动测试
/// [已废弃] Worker 进程(新架构使用 daemon 统一 tokio task
///
/// 示例:
/// ias worker --sock /tmp/ias_daemon.sock
+9
View File
@@ -2,6 +2,7 @@ use crate::llm::types::Message;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use uuid::Uuid;
/// 摘要原因
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
@@ -44,6 +45,8 @@ pub struct ChatSession {
pub token_budget: u32,
/// 12 小时空闲阈值(秒)
pub idle_timeout_secs: i64,
/// 会话 ID(用于追踪工具调用往返)
pub session_id: Uuid,
}
impl Default for ChatSession {
@@ -58,6 +61,7 @@ impl Default for ChatSession {
last_user_at: None,
token_budget: 28000,
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));
}
/// 添加一条系统说明消息(如已批准的工具)
pub fn add_system_note(&mut self, note: &str) {
self.messages.push(Message::system(note));
}
/// 是否因空闲超过阈值需要摘要
pub fn is_idle_timeout(&self) -> bool {
if let Some(last) = self.last_user_at {
+491 -716
View File
File diff suppressed because it is too large Load Diff
+1
View File
@@ -149,6 +149,7 @@ pub async fn recv_task(stream: &mut UnixStream) -> Result<TaskFrame, String> {
}
/// 接收并解析为 OutputFrame
#[allow(dead_code)]
pub async fn recv_output(stream: &mut UnixStream) -> Result<OutputFrame, String> {
let v = recv_frame(stream).await?;
serde_json::from_value(v).map_err(|e| format!("解析 OutputFrame 失败: {}", e))
+58 -5
View File
@@ -4,6 +4,7 @@ use crate::llm::provider::{BoxedProvider, StreamReceiver, create_provider};
use crate::llm::types::{ConversationConfig, Message, StreamChunk, ToolCall, Usage};
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
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,
>;
/// `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 {
config: ConversationConfig,
provider: BoxedProvider,
session: Arc<Mutex<ChatSession>>,
tool_executor: Option<ToolExecutor>,
/// 是否有异步工具调用等待结果返回
pending_async_tools: Arc<AtomicBool>,
}
impl Conversation {
@@ -33,12 +49,18 @@ impl Conversation {
provider,
session: Arc::new(Mutex::new(ChatSession::new())),
tool_executor: None,
pending_async_tools: Arc::new(AtomicBool::new(false)),
})
}
pub fn set_tool_executor(&mut self, executor: ToolExecutor) {
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>> {
Arc::clone(&self.session)
}
@@ -113,10 +135,7 @@ impl Conversation {
}
/// 对话 + 工具循环(带上下文管理)
pub async fn chat_with_tools(
&self,
user_message: String,
) -> Result<(String, bool, Option<Usage>), String> {
pub async fn chat_with_tools(&self, user_message: String) -> Result<ChatResult, String> {
// ── 空闲超时检查 ──
{
let s = self.session.lock().await;
@@ -128,7 +147,17 @@ impl Conversation {
}
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 last_usage = None;
let mut turn_count = 0u32;
@@ -185,6 +214,7 @@ impl Conversation {
.await
.add_assistant_tool_calls(calls.clone());
let mut has_async = false;
for tc in &calls {
let result = match &self.tool_executor {
Some(exec) => {
@@ -195,6 +225,14 @@ impl Conversation {
}
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);
self.session.lock().await.add_tool_result(
&tc.id,
@@ -202,6 +240,16 @@ impl Conversation {
&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;
}
}
@@ -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
View File
@@ -3,5 +3,5 @@ pub mod deepseek;
pub mod provider;
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};
+7 -3
View File
@@ -1,3 +1,4 @@
mod channel;
mod cli;
mod context;
mod daemon;
@@ -5,6 +6,7 @@ mod db;
mod ipc;
mod llm;
mod logger;
mod queue;
mod scheduler;
mod state;
mod tools;
@@ -15,11 +17,11 @@ use clap::Parser;
use cli::{AmapAction, Cli, Commands, MemosAction, ToolCommand};
use context::MemoryStore;
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 tools::approval::{ApprovalDecision, ApprovalManager};
use tools::types::ExecutionContext;
use tracing::{error, info};
use tracing::{error, info, warn};
use wechat::client::WeChatClient;
#[tokio::main]
@@ -105,6 +107,8 @@ async fn main() {
cmd_daemon(sock, &database, &file_state, &memory_store).await;
}
Commands::Worker { sock } => {
warn!("Worker 模式已废弃,请使用 daemon 模式");
#[allow(deprecated)]
if let Err(e) = worker::run(&sock).await {
error!("Worker 失败: {}", e);
}
@@ -665,7 +669,7 @@ async fn generate_reply_with_tools(
user_text: String,
db: &Option<Arc<Database>>,
) -> 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 {
info!(
"📊 Token: {} 总 / {} 提示 / {} 生成 | 缓存命中 {} | 缓存未命中 {}",
+552
View File
@@ -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,
},
}
}
}
// ─── 公平轮转队列 ───
/// 多渠道公平轮转消息队列
///
/// 按 ChannelIdplatform + 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();
}
}
+5
View File
@@ -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};
+212
View File
@@ -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!("队列系统已完全关闭");
}
+364
View File
@@ -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");
}
}
+934
View File
@@ -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.nameURL 编码)
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 = &param_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
View File
@@ -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
//! 2. 读取 TaskFrame (用户消息 + 上下文 + 环境变量)
//! 3. 注入环境变量,构建 Conversation,执行 LLM 对话 + 工具调用
@@ -103,7 +106,8 @@ pub async fn run(sock_path: &str) -> Result<(), String> {
conv.set_tool_executor(executor);
// 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())
.await
.map_err(|e| format!("LLM 对话失败: {}", e))?;