583 lines
19 KiB
Rust
583 lines
19 KiB
Rust
//! ## 多渠道公平轮转消息队列 (MessageQueue)
|
||
//!
|
||
//! ### 设计目标
|
||
//!
|
||
//! 多用户场景下保证**公平性**：如果一个用户连续发了很多消息，不会让其他用户饿死。
|
||
//! 出队顺序是**轮转（round-robin）**的：
|
||
//! - A 发了 3 条、B 发了 1 条 → 处理顺序 A1, B1, A2, A3
|
||
//!
|
||
//! ### 消息关联追踪
|
||
//!
|
||
//! 每条消息有一个 `correlation_id`（Uuid），用来追踪一次用户请求的完整生命周期：
|
||
//! 用户消息 → LLM 处理 → 工具调用 → 工具结果回喂 → 最终回复发送。
|
||
//! 这样即使多用户并发，回复也不会串。
|
||
//!
|
||
//! ### 架构位置
|
||
//!
|
||
//! ```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>,
|
||
/// 该工具调用是否已被用户批准（跳过审批流程）
|
||
approved: bool,
|
||
},
|
||
/// 工具执行结果 → 回喂 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,
|
||
approved: false,
|
||
},
|
||
}
|
||
}
|
||
|
||
/// 快速构造一条已批准的工具调用（跳过审批）
|
||
pub fn approved_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,
|
||
approved: true,
|
||
},
|
||
}
|
||
}
|
||
|
||
/// 快速构造一条工具结果
|
||
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)
|
||
&& 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).is_some_and(|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();
|
||
}
|
||
}
|