feat: 拆分 worker 架构并增强 assistant 安全保护

- 新增 daemon/channel/assistant/function/scheduler/logging worker 架构和 JetStream 消息协议

- 修复邮件通知误入 assistant、迁移 checksum、日志 subject 和工具调用无限循环风险

- 新增 assistant 工具轮次熔断、工具结果截断和提示词软收敛规则

- 补充邮件推送关键字、可信发件人配置、升级日志和验证测试
This commit is contained in:
2026-07-10 00:41:36 +08:00
parent 7deae286ca
commit af6cfdaa83
70 changed files with 5953 additions and 1608 deletions
+13 -1
View File
@@ -1,15 +1,27 @@
[package]
name = "ias-common"
version = "0.1.2"
version = "0.2.0"
edition = "2024"
[dependencies]
# ── NATS JetStream ──
async-nats = "0.49.1"
# ── HTTP 客户端 ──
reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "rustls-tls"] } # HTTP 请求(微信 API + DeepSeek API + 工具调用)
# ── 序列化 ──
serde = { version = "1.0", features = ["derive"] } # 序列化/反序列化
serde_json = "1.0" # JSON 处理
serde_yaml = "0.9"
# ── 异步流工具 ──
futures-util = "0.3"
tokio = { version = "1.0", features = ["rt", "sync"] }
# ── 错误处理 ──
anyhow = "1"
thiserror = "2"
# ── 时间处理 ──
chrono = { version = "0.4", features = ["serde"] }
# ── 唯一标识 ──
uuid = { version = "1", features = ["v4", "serde"] }
[dev-dependencies]
tokio = { version = "1.0", features = ["rt-multi-thread", "macros", "sync", "time"] }
+123
View File
@@ -0,0 +1,123 @@
use std::path::PathBuf;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone)]
pub struct RuntimePaths {
pub dir: PathBuf,
pub pid_file: PathBuf,
pub socket_file: PathBuf,
pub state_file: PathBuf,
pub log_file: PathBuf,
}
impl RuntimePaths {
pub fn new() -> Self {
let dir = std::env::var_os("IAS_RUNTIME_DIR")
.map(PathBuf::from)
.or_else(|| {
std::env::var_os("XDG_RUNTIME_DIR")
.map(|runtime_dir| PathBuf::from(runtime_dir).join("ias"))
})
.unwrap_or_else(|| {
let user = std::env::var("USER").unwrap_or_else(|_| "user".to_string());
std::env::temp_dir().join(format!("ias-{user}"))
});
Self {
pid_file: dir.join("ias.pid"),
socket_file: std::env::var_os("IAS_CONTROL_SOCKET")
.map(PathBuf::from)
.unwrap_or_else(|| dir.join("ias.sock")),
state_file: dir.join("ias.state.json"),
log_file: dir.join("ias.log"),
dir,
}
}
pub fn ensure_dir(&self) -> anyhow::Result<()> {
std::fs::create_dir_all(&self.dir)?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(&self.dir, std::fs::Permissions::from_mode(0o700))?;
}
Ok(())
}
}
impl Default for RuntimePaths {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ServicePhase {
Starting,
Running,
Degraded,
Stopping,
Stopped,
Failed,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum WorkerPhase {
Starting,
Running,
Stopped,
Failed,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkerStatus {
pub name: String,
pub role: String,
pub pid: Option<u32>,
pub phase: WorkerPhase,
pub last_heartbeat_at: Option<DateTime<Utc>>,
pub last_error: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamStatus {
pub name: String,
pub subjects: Vec<String>,
pub messages: u64,
pub bytes: u64,
pub consumer_lag: Option<u64>,
pub dlq_messages: Option<u64>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MqStatus {
pub connected: bool,
pub nats_url: String,
pub stream_prefix: String,
pub streams: Vec<StreamStatus>,
pub last_error: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RuntimeSnapshot {
pub phase: ServicePhase,
pub pid: u32,
pub started_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
pub http_addr: Option<String>,
pub mq: MqStatus,
pub workers: Vec<WorkerStatus>,
pub last_error: Option<String>,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ControlResponse {
Status { snapshot: RuntimeSnapshot },
Ack { message: String },
Error { message: String },
}
+2
View File
@@ -1,3 +1,5 @@
pub mod control;
pub mod error;
pub mod model;
pub mod mq;
pub mod queue;
+460
View File
@@ -0,0 +1,460 @@
use std::future::Future;
use std::sync::Arc;
use std::time::Duration;
use crate::control::StreamStatus;
use async_nats::jetstream::consumer::{self, PullConsumer};
use async_nats::jetstream::message::PublishMessage;
use async_nats::jetstream::{self, stream};
use futures_util::{FutureExt, StreamExt};
use serde::Serialize;
use serde::de::DeserializeOwned;
use serde_json::Value;
use tokio::sync::Semaphore;
use tokio::task::JoinSet;
use uuid::Uuid;
use super::message::{DeadLetterPayload, MessageEnvelope, now_millis};
use super::subjects::{StreamKind, StreamNames, dlq_subject};
#[derive(Debug, Clone)]
pub struct JetStreamConfig {
pub nats_url: String,
pub stream_prefix: String,
pub ack_wait_secs: u64,
pub max_deliver: u64,
pub consumer_concurrency: usize,
pub log_retention_days: u64,
pub log_max_bytes: i64,
}
impl JetStreamConfig {
pub fn from_env() -> Self {
Self {
nats_url: std::env::var("NATS_URL")
.ok()
.filter(|value| !value.trim().is_empty())
.unwrap_or_else(|| "nats://127.0.0.1:4222".to_string()),
stream_prefix: std::env::var("IAS_NATS_STREAM_PREFIX")
.ok()
.filter(|value| !value.trim().is_empty())
.unwrap_or_else(|| "IA".to_string()),
ack_wait_secs: read_u64_env("IAS_NATS_ACK_WAIT_SECS", 60),
max_deliver: read_u64_env("IAS_NATS_MAX_DELIVER", 5),
consumer_concurrency: read_usize_env("IAS_NATS_CONSUMER_CONCURRENCY", 8),
log_retention_days: read_u64_env("IAS_LOG_EVENT_RETENTION_DAYS", 30),
log_max_bytes: read_i64_env("IAS_LOG_EVENT_MAX_BYTES", 512 * 1024 * 1024),
}
}
pub fn names(&self) -> StreamNames {
StreamNames::new(self.stream_prefix.clone())
}
pub fn ack_wait(&self) -> Duration {
Duration::from_secs(self.ack_wait_secs)
}
}
#[derive(Debug, Clone)]
pub struct StreamSpec {
pub kind: StreamKind,
pub name: String,
pub subjects: Vec<String>,
pub max_age: Option<Duration>,
pub max_bytes: Option<i64>,
}
#[derive(Debug, Clone)]
pub struct ConsumerSpec {
pub stream: String,
pub durable_name: String,
pub filter_subject: String,
pub ack_wait: Duration,
pub max_deliver: u64,
}
pub fn canonical_stream_specs(config: &JetStreamConfig) -> Vec<StreamSpec> {
let names = config.names();
vec![
StreamSpec {
kind: StreamKind::Inbound,
name: names.name(StreamKind::Inbound),
subjects: vec!["inbound.wechat.*".to_string(), "inbound.mail.*".to_string()],
max_age: None,
max_bytes: None,
},
StreamSpec {
kind: StreamKind::Function,
name: names.name(StreamKind::Function),
subjects: vec![
"function.call.*".to_string(),
"function.result.*".to_string(),
],
max_age: None,
max_bytes: None,
},
StreamSpec {
kind: StreamKind::Outbound,
name: names.name(StreamKind::Outbound),
subjects: vec![
"outbound.wechat.*".to_string(),
"outbound.mail.*".to_string(),
],
max_age: None,
max_bytes: None,
},
StreamSpec {
kind: StreamKind::Schedule,
name: names.name(StreamKind::Schedule),
subjects: vec!["schedule.trigger.*".to_string()],
max_age: None,
max_bytes: None,
},
StreamSpec {
kind: StreamKind::Logs,
name: names.name(StreamKind::Logs),
subjects: vec![
"logs.daemon.*".to_string(),
"logs.worker.*".to_string(),
"logs.audit.*".to_string(),
"logs.http.*".to_string(),
],
max_age: Some(Duration::from_secs(
config.log_retention_days * 24 * 60 * 60,
)),
max_bytes: Some(config.log_max_bytes),
},
StreamSpec {
kind: StreamKind::Dlq,
name: names.name(StreamKind::Dlq),
subjects: vec!["dlq.*".to_string()],
max_age: None,
max_bytes: None,
},
]
}
#[derive(Clone)]
pub struct JetStreamBus {
jetstream: jetstream::Context,
config: JetStreamConfig,
}
impl JetStreamBus {
pub async fn connect(config: JetStreamConfig) -> anyhow::Result<Self> {
let client = async_nats::connect(config.nats_url.clone()).await?;
let jetstream = jetstream::new(client);
Ok(Self { jetstream, config })
}
pub fn config(&self) -> &JetStreamConfig {
&self.config
}
pub async fn ensure_streams(&self) -> anyhow::Result<()> {
for spec in canonical_stream_specs(&self.config) {
let mut stream_config = stream::Config {
name: spec.name,
subjects: spec.subjects,
storage: stream::StorageType::File,
..Default::default()
};
if let Some(max_age) = spec.max_age {
stream_config.max_age = max_age;
}
if let Some(max_bytes) = spec.max_bytes {
stream_config.max_bytes = max_bytes;
}
self.jetstream.get_or_create_stream(stream_config).await?;
}
Ok(())
}
pub async fn stream_statuses(&self) -> anyhow::Result<Vec<StreamStatus>> {
let mut statuses = Vec::new();
for spec in canonical_stream_specs(&self.config) {
let mut stream = self.jetstream.get_stream(spec.name.clone()).await?;
let info = stream.info().await?;
statuses.push(StreamStatus {
name: info.config.name.clone(),
subjects: info.config.subjects.clone(),
messages: info.state.messages,
bytes: info.state.bytes,
consumer_lag: None,
dlq_messages: None,
});
}
Ok(statuses)
}
pub async fn publish_envelope<T>(
&self,
subject: &str,
envelope: &MessageEnvelope<T>,
) -> anyhow::Result<()>
where
T: Serialize,
{
let payload = serde_json::to_vec(envelope)?;
let publish = PublishMessage::build()
.payload(payload.into())
.message_id(&envelope.message_id);
let ack = self
.jetstream
.send_publish(subject.to_string(), publish)
.await?;
ack.await?;
Ok(())
}
pub async fn pull_consumer(&self, spec: ConsumerSpec) -> anyhow::Result<PullConsumer> {
let stream = self.jetstream.get_stream(spec.stream).await?;
let consumer = stream
.get_or_create_consumer(
&spec.durable_name,
consumer::pull::Config {
durable_name: Some(spec.durable_name.clone()),
filter_subject: spec.filter_subject,
ack_policy: consumer::AckPolicy::Explicit,
ack_wait: spec.ack_wait,
max_deliver: spec.max_deliver as i64,
..Default::default()
},
)
.await?;
Ok(consumer)
}
/// 消费 JetStream 消息,按 `consumer_concurrency` 并发执行 handler。
///
/// 语义:
/// - handler 返回 `Ok`ack 消息。
/// - handler 返回 `Err` 且未达到 `max_deliver`:显式 `nak` 触发快速重投。
/// - handler 返回 `Err` 且已达到 `max_deliver`:转入 DLQ 并 ack。
/// - 反序列化失败:立即转入 DLQ 并 ack(永久错误)。
pub async fn consume_envelopes<T, F, Fut>(
&self,
spec: ConsumerSpec,
handler: F,
) -> anyhow::Result<()>
where
T: DeserializeOwned + Serialize + Send + 'static,
F: Fn(MessageEnvelope<T>) -> Fut + Send + Sync + 'static,
Fut: Future<Output = anyhow::Result<()>> + Send + 'static,
{
let consumer = self.pull_consumer(spec.clone()).await?;
let mut messages = consumer.messages().await?;
let concurrency = self.config.consumer_concurrency.max(1);
let semaphore = Arc::new(Semaphore::new(concurrency));
let handler = Arc::new(handler);
let mut join_set: JoinSet<anyhow::Result<()>> = JoinSet::new();
loop {
drain_ready_consumer_tasks(&mut join_set);
tokio::select! {
msg = messages.next() => {
let Some(message) = msg else { break; };
let message = message?;
let permit = semaphore.clone().acquire_owned().await?;
let handler = handler.clone();
let bus = self.clone();
let spec = spec.clone();
join_set.spawn(async move {
let _permit = permit;
bus.process_message(&spec, message, &*handler).await
});
}
result = join_set.join_next(), if !join_set.is_empty() => {
log_consumer_task_result(result);
}
}
}
while let Some(result) = join_set.join_next().await {
log_consumer_task_result(Some(result));
}
Ok(())
}
async fn process_message<T, F, Fut>(
&self,
spec: &ConsumerSpec,
message: jetstream::Message,
handler: &F,
) -> anyhow::Result<()>
where
T: DeserializeOwned + Serialize + Send + 'static,
F: Fn(MessageEnvelope<T>) -> Fut + Send + Sync + 'static,
Fut: Future<Output = anyhow::Result<()>> + Send + 'static,
{
let delivered = message.info().ok().map(|info| info.delivered);
let reached_limit = delivered.is_some_and(|value| value >= spec.max_deliver as i64);
let envelope = match serde_json::from_slice::<MessageEnvelope<T>>(&message.payload) {
Ok(envelope) => envelope.with_attempt_hint(delivered.unwrap_or(1) as u32),
Err(error) => {
self.publish_dlq(
spec,
delivered,
&message.payload,
format!("消息反序列化失败: {error}"),
)
.await?;
ack_message(&message).await?;
return Ok(());
}
};
match handler(envelope).await {
Ok(()) => {
ack_message(&message).await?;
}
Err(error) => {
eprintln!(
"处理 JetStream 消息失败 durable={} filter={}: {error:#}",
spec.durable_name, spec.filter_subject
);
if reached_limit {
self.publish_dlq(
spec,
delivered,
&message.payload,
format!("处理失败达到最大投递次数 {}: {error}", spec.max_deliver),
)
.await?;
ack_message(&message).await?;
} else {
let _ = message
.ack_with(async_nats::jetstream::AckKind::Nak(None))
.await;
}
}
}
Ok(())
}
async fn publish_dlq(
&self,
spec: &ConsumerSpec,
delivered: Option<i64>,
original_payload: &[u8],
error: String,
) -> anyhow::Result<()> {
let payload: Value = serde_json::from_slice(original_payload)
.unwrap_or_else(|_| Value::String("<unparseable>".to_string()));
let original_message_id = payload
.get("message_id")
.and_then(Value::as_str)
.map(ToOwned::to_owned);
// 用原始 message_id 作为 DLQ 幂等键,缺失时回退 uuid,避免同毫秒并发碰撞丢消息。
let dlq_message_id = match &original_message_id {
Some(message_id) => format!("dlq/{message_id}"),
None => format!("dlq/{}", Uuid::new_v4()),
};
let trace_id = payload
.get("trace_id")
.and_then(Value::as_str)
.map(ToOwned::to_owned)
.unwrap_or_else(|| format!("dlq-{}", now_millis()));
let dlq = DeadLetterPayload {
original_subject: spec.filter_subject.clone(),
original_stream: Some(spec.stream.clone()),
original_consumer: Some(spec.durable_name.clone()),
original_message_id,
trace_id: Some(trace_id.clone()),
delivered,
error,
payload,
created_at_ms: now_millis(),
};
let subject = dlq_subject(&spec.stream);
let envelope = MessageEnvelope::new(subject.as_str(), dlq_message_id, trace_id, dlq);
self.publish_envelope(subject.as_str(), &envelope).await
}
}
fn read_u64_env(name: &str, default: u64) -> u64 {
std::env::var(name)
.ok()
.and_then(|value| value.trim().parse::<u64>().ok())
.filter(|value| *value > 0)
.unwrap_or(default)
}
fn read_usize_env(name: &str, default: usize) -> usize {
std::env::var(name)
.ok()
.and_then(|value| value.trim().parse::<usize>().ok())
.filter(|value| *value > 0)
.unwrap_or(default)
}
fn read_i64_env(name: &str, default: i64) -> i64 {
std::env::var(name)
.ok()
.and_then(|value| value.trim().parse::<i64>().ok())
.filter(|value| *value > 0)
.unwrap_or(default)
}
async fn ack_message(message: &jetstream::Message) -> anyhow::Result<()> {
message
.ack()
.await
.map_err(|error| anyhow::anyhow!(error.to_string()))
}
fn log_consumer_task_result(result: Option<Result<anyhow::Result<()>, tokio::task::JoinError>>) {
match result {
Some(Ok(Ok(()))) | None => {}
Some(Ok(Err(error))) => eprintln!("消费者任务处理失败: {error:#}"),
Some(Err(error)) => eprintln!("消费者任务异常退出: {error}"),
}
}
fn drain_ready_consumer_tasks(join_set: &mut JoinSet<anyhow::Result<()>>) {
while let Some(result) = join_set.join_next().now_or_never().flatten() {
log_consumer_task_result(Some(result));
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn canonical_specs_cover_required_streams() {
let config = JetStreamConfig {
nats_url: "nats://127.0.0.1:4222".to_string(),
stream_prefix: "TEST".to_string(),
ack_wait_secs: 60,
max_deliver: 5,
consumer_concurrency: 8,
log_retention_days: 30,
log_max_bytes: 1024,
};
let streams = canonical_stream_specs(&config);
let names = streams
.iter()
.map(|stream| stream.name.as_str())
.collect::<Vec<_>>();
assert_eq!(
names,
vec![
"TEST_INBOUND",
"TEST_FUNCTION",
"TEST_OUTBOUND",
"TEST_SCHEDULE",
"TEST_LOGS",
"TEST_DLQ",
]
);
assert!(
streams
.iter()
.any(|stream| stream.subjects.contains(&"outbound.wechat.*".to_string()))
);
}
}
+279
View File
@@ -0,0 +1,279 @@
use chrono::Utc;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use uuid::Uuid;
use crate::queue::ChannelMessage;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct MessageEnvelope<T> {
pub message_id: String,
pub topic: String,
pub trace_id: String,
pub created_at_ms: i64,
pub attempt_hint: u32,
pub payload: T,
}
impl<T> MessageEnvelope<T> {
pub fn new(
topic: impl Into<String>,
message_id: impl Into<String>,
trace_id: impl Into<String>,
payload: T,
) -> Self {
Self {
message_id: message_id.into(),
topic: topic.into(),
trace_id: trace_id.into(),
created_at_ms: now_millis(),
attempt_hint: 0,
payload,
}
}
pub fn with_attempt_hint(mut self, attempt_hint: u32) -> Self {
self.attempt_hint = attempt_hint;
self
}
}
pub fn now_millis() -> i64 {
Utc::now().timestamp_millis()
}
pub fn new_trace_id() -> String {
Uuid::new_v4().to_string()
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct InboundPayload {
pub channel: String,
pub account_id: String,
pub from_user_id: String,
pub content: String,
pub context: Option<String>,
pub external_message_id: Option<String>,
pub session_id: Option<String>,
pub group_id: Option<String>,
pub create_time_ms: Option<i64>,
}
impl InboundPayload {
pub fn idempotency_key(&self) -> String {
let external_id = self
.external_message_id
.as_deref()
.filter(|value| !value.trim().is_empty())
.unwrap_or("no_external_id");
format!("{}/{}/{}", self.channel, self.account_id, external_id)
}
pub fn into_channel_message(self) -> ChannelMessage {
ChannelMessage {
channel: self.channel,
account_id: self.account_id,
from_user_id: self.from_user_id,
content: self.content,
context: self.context,
turn_id: None,
external_message_id: self.external_message_id,
session_id: self.session_id,
group_id: self.group_id,
create_time_ms: self.create_time_ms,
}
}
}
impl From<ChannelMessage> for InboundPayload {
fn from(message: ChannelMessage) -> Self {
Self {
channel: message.channel,
account_id: message.account_id,
from_user_id: message.from_user_id,
content: message.content,
context: message.context,
external_message_id: message.external_message_id,
session_id: message.session_id,
group_id: message.group_id,
create_time_ms: message.create_time_ms,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum OutboundCommandKind {
Text,
TypingStart,
TypingRefresh,
TypingStop,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct OutboundCommandPayload {
pub command_id: String,
pub kind: OutboundCommandKind,
pub channel: String,
pub account_id: String,
pub to_user_id: String,
pub content: Option<String>,
pub context: Option<String>,
pub client_id: Option<String>,
pub turn_id: Option<String>,
}
impl OutboundCommandPayload {
pub fn text(
channel: impl Into<String>,
account_id: impl Into<String>,
to_user_id: impl Into<String>,
content: impl Into<String>,
context: Option<String>,
turn_id: Option<String>,
) -> Self {
Self {
command_id: Uuid::new_v4().to_string(),
kind: OutboundCommandKind::Text,
channel: channel.into(),
account_id: account_id.into(),
to_user_id: to_user_id.into(),
content: Some(content.into()),
context,
client_id: None,
turn_id,
}
}
pub fn typing(
kind: OutboundCommandKind,
channel: impl Into<String>,
account_id: impl Into<String>,
to_user_id: impl Into<String>,
context: Option<String>,
client_id: Option<String>,
turn_id: Option<String>,
) -> Self {
Self {
command_id: Uuid::new_v4().to_string(),
kind,
channel: channel.into(),
account_id: account_id.into(),
to_user_id: to_user_id.into(),
content: None,
context,
client_id,
turn_id,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FunctionCallPayload {
pub call_id: String,
pub name: String,
pub arguments: String,
pub channel: ChannelMessage,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FunctionResultPayload {
pub call_id: String,
pub name: String,
pub result: String,
pub channel: ChannelMessage,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ScheduledTriggerPayload {
pub fire_id: String,
pub schedule_id: i64,
pub schedule_type: String,
pub target_topic: String,
pub fire_time_ms: i64,
pub payload: Value,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct LogEventPayload {
pub event_id: String,
pub level: String,
pub source: String,
pub message: String,
pub trace_id: String,
pub message_id: Option<String>,
pub command_id: Option<String>,
pub call_id: Option<String>,
pub fire_id: Option<String>,
pub fields: Value,
pub created_at_ms: i64,
}
impl LogEventPayload {
pub fn new(
level: impl Into<String>,
source: impl Into<String>,
message: impl Into<String>,
trace_id: impl Into<String>,
fields: Value,
) -> Self {
Self {
event_id: Uuid::new_v4().to_string(),
level: level.into(),
source: source.into(),
message: message.into(),
trace_id: trace_id.into(),
message_id: None,
command_id: None,
call_id: None,
fire_id: None,
fields,
created_at_ms: now_millis(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct DeadLetterPayload {
pub original_subject: String,
pub original_stream: Option<String>,
pub original_consumer: Option<String>,
pub original_message_id: Option<String>,
pub trace_id: Option<String>,
pub delivered: Option<i64>,
pub error: String,
pub payload: Value,
pub created_at_ms: i64,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn inbound_idempotency_prefers_external_message_id() {
let payload = InboundPayload {
channel: "wechat".to_string(),
account_id: "acct".to_string(),
from_user_id: "user".to_string(),
content: "hello".to_string(),
context: None,
external_message_id: Some("msg-1".to_string()),
session_id: None,
group_id: None,
create_time_ms: None,
};
assert_eq!(payload.idempotency_key(), "wechat/acct/msg-1");
}
#[test]
fn envelope_preserves_stable_ids() {
let payload = serde_json::json!({"ok": true});
let envelope = MessageEnvelope::new("inbound.wechat.acct", "msg-1", "trace-1", payload);
assert_eq!(envelope.message_id, "msg-1");
assert_eq!(envelope.trace_id, "trace-1");
assert_eq!(envelope.attempt_hint, 0);
}
}
+14
View File
@@ -0,0 +1,14 @@
mod bus;
mod message;
mod subjects;
pub use bus::{ConsumerSpec, JetStreamBus, JetStreamConfig, StreamSpec, canonical_stream_specs};
pub use message::{
DeadLetterPayload, FunctionCallPayload, FunctionResultPayload, InboundPayload, LogEventPayload,
MessageEnvelope, OutboundCommandKind, OutboundCommandPayload, ScheduledTriggerPayload,
new_trace_id, now_millis,
};
pub use subjects::{
StreamKind, StreamNames, Subject, dlq_subject, function_call_subject, function_result_subject,
inbound_subject, log_subject, outbound_subject, schedule_trigger_subject,
};
+163
View File
@@ -0,0 +1,163 @@
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StreamKind {
Inbound,
Function,
Outbound,
Schedule,
Logs,
Dlq,
}
impl StreamKind {
pub fn suffix(self) -> &'static str {
match self {
StreamKind::Inbound => "INBOUND",
StreamKind::Function => "FUNCTION",
StreamKind::Outbound => "OUTBOUND",
StreamKind::Schedule => "SCHEDULE",
StreamKind::Logs => "LOGS",
StreamKind::Dlq => "DLQ",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StreamNames {
prefix: String,
}
impl StreamNames {
pub fn new(prefix: impl Into<String>) -> Self {
Self {
prefix: prefix.into(),
}
}
pub fn from_env() -> Self {
Self::new(
std::env::var("IAS_NATS_STREAM_PREFIX")
.ok()
.filter(|value| !value.trim().is_empty())
.unwrap_or_else(|| "IA".to_string()),
)
}
pub fn name(&self, kind: StreamKind) -> String {
format!("{}_{}", self.prefix, kind.suffix())
}
pub fn prefix(&self) -> &str {
&self.prefix
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Subject(String);
impl Subject {
pub fn as_str(&self) -> &str {
&self.0
}
}
impl From<Subject> for String {
fn from(value: Subject) -> Self {
value.0
}
}
impl std::fmt::Display for Subject {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str(&self.0)
}
}
pub fn inbound_subject(channel: &str, account_id: &str) -> Subject {
Subject(format!(
"inbound.{}.{}",
subject_token(channel),
subject_token(account_id)
))
}
pub fn outbound_subject(channel: &str, account_id: &str) -> Subject {
Subject(format!(
"outbound.{}.{}",
subject_token(channel),
subject_token(account_id)
))
}
pub fn function_call_subject(name: &str) -> Subject {
Subject(format!("function.call.{}", subject_token(name)))
}
pub fn function_result_subject(call_id: &str) -> Subject {
Subject(format!("function.result.{}", subject_token(call_id)))
}
pub fn schedule_trigger_subject(target_topic: &str) -> Subject {
Subject(format!("schedule.trigger.{}", subject_token(target_topic)))
}
pub fn log_subject(source: &str) -> Subject {
Subject(format!("logs.{}", subject_path(source)))
}
pub fn dlq_subject(original_stream: &str) -> Subject {
Subject(format!("dlq.{}", subject_token(original_stream)))
}
fn subject_token(value: &str) -> String {
let token = value
.chars()
.map(|ch| {
if ch.is_ascii_alphanumeric() || ch == '_' || ch == '-' {
ch
} else {
'_'
}
})
.collect::<String>();
if token.trim_matches('_').is_empty() {
"unknown".to_string()
} else {
token
}
}
fn subject_path(value: &str) -> String {
value
.split('.')
.map(subject_token)
.collect::<Vec<_>>()
.join(".")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn stream_names_use_configured_prefix() {
let names = StreamNames::new("TEST");
assert_eq!(names.name(StreamKind::Inbound), "TEST_INBOUND");
assert_eq!(names.name(StreamKind::Dlq), "TEST_DLQ");
}
#[test]
fn subject_tokens_are_nats_safe() {
let subject = inbound_subject("we/chat", "acct.1");
assert_eq!(subject.as_str(), "inbound.we_chat.acct_1");
}
#[test]
fn log_subject_preserves_hierarchical_source_segments() {
let subject = log_subject("worker.assistant");
assert_eq!(subject.as_str(), "logs.worker.assistant");
}
}
+1 -1
View File
@@ -7,7 +7,7 @@ use crate::model::ai::ChatCompletionToolCall;
pub enum QueueMessage {
Channel(ChannelMessage),
Notice(ChannelMessage, String),
Aissitant(ChannelMessage, AssistantMessage),
Assistant(ChannelMessage, AssistantMessage),
Tool(ChannelMessage, String, String, String),
Tools(ChannelMessage, Vec<ToolMessage>),
}
+256
View File
@@ -0,0 +1,256 @@
use ias_common::mq::{
ConsumerSpec, DeadLetterPayload, InboundPayload, JetStreamBus, JetStreamConfig,
MessageEnvelope, StreamKind, StreamNames, inbound_subject,
};
#[tokio::test]
#[ignore = "requires a running NATS JetStream server from docker-compose.nats.yml"]
async fn jetstream_publish_and_ack_roundtrip() -> anyhow::Result<()> {
let config = JetStreamConfig::from_env();
let bus = JetStreamBus::connect(config.clone()).await?;
bus.ensure_streams().await?;
let subject = inbound_subject("wechat", "integration");
let payload = InboundPayload {
channel: "wechat".to_string(),
account_id: "integration".to_string(),
from_user_id: "tester".to_string(),
content: "ping".to_string(),
context: None,
external_message_id: Some("integration-msg-1".to_string()),
session_id: None,
group_id: None,
create_time_ms: None,
};
let envelope = MessageEnvelope::new(
subject.as_str(),
payload.idempotency_key(),
"trace-integration",
payload,
);
bus.publish_envelope(subject.as_str(), &envelope).await?;
let names = StreamNames::new(config.stream_prefix.clone());
let spec = ConsumerSpec {
stream: names.name(StreamKind::Inbound),
durable_name: "integration-ack-test".to_string(),
filter_subject: subject.to_string(),
ack_wait: config.ack_wait(),
max_deliver: config.max_deliver,
};
let (tx, mut rx) = tokio::sync::mpsc::channel(1);
let consumer_bus = bus.clone();
let handle = tokio::spawn(async move {
consumer_bus
.consume_envelopes::<InboundPayload, _, _>(spec, move |message| {
let tx = tx.clone();
async move {
tx.send(message.message_id).await?;
Ok(())
}
})
.await
});
let message_id = tokio::time::timeout(std::time::Duration::from_secs(10), rx.recv())
.await?
.ok_or_else(|| anyhow::anyhow!("consumer closed before receiving test message"))?;
handle.abort();
assert_eq!(message_id, "wechat/integration/integration-msg-1");
Ok(())
}
#[tokio::test]
#[ignore = "requires a running NATS JetStream server from docker-compose.nats.yml"]
async fn jetstream_redelivery_reaches_dlq_after_max_deliver() -> anyhow::Result<()> {
let mut config = JetStreamConfig::from_env();
config.ack_wait_secs = 1;
config.max_deliver = 2;
let bus = JetStreamBus::connect(config.clone()).await?;
bus.ensure_streams().await?;
let names = StreamNames::new(config.stream_prefix.clone());
let unique = uuid::Uuid::new_v4().to_string();
let account_id = format!("dlq-{unique}");
let subject = inbound_subject("wechat", &account_id);
let payload = InboundPayload {
channel: "wechat".to_string(),
account_id: account_id.clone(),
from_user_id: "tester".to_string(),
content: "force retry".to_string(),
context: None,
external_message_id: Some(unique.clone()),
session_id: None,
group_id: None,
create_time_ms: None,
};
let message_id = format!("wechat/{account_id}/{unique}");
let envelope = MessageEnvelope::new(
subject.as_str(),
message_id.clone(),
format!("trace-{unique}"),
payload,
);
let (dlq_tx, mut dlq_rx) = tokio::sync::mpsc::channel(1);
let dlq_bus = bus.clone();
let dlq_stream = names.name(StreamKind::Dlq);
let dlq_ack_wait = config.ack_wait();
let dlq_max_deliver = config.max_deliver;
let dlq_handle = tokio::spawn(async move {
dlq_bus
.consume_envelopes::<DeadLetterPayload, _, _>(
ConsumerSpec {
stream: dlq_stream,
durable_name: format!("integration-dlq-{unique}"),
filter_subject: "dlq.>".to_string(),
ack_wait: dlq_ack_wait,
max_deliver: dlq_max_deliver,
},
move |message| {
let dlq_tx = dlq_tx.clone();
let expected = message_id.clone();
async move {
if message.payload.original_message_id.as_deref() == Some(&expected) {
dlq_tx.send(message.payload.original_message_id).await?;
}
Ok(())
}
},
)
.await
});
let retry_bus = bus.clone();
let retry_spec = ConsumerSpec {
stream: StreamNames::new(config.stream_prefix.clone()).name(StreamKind::Inbound),
durable_name: format!("integration-redelivery-{account_id}"),
filter_subject: subject.to_string(),
ack_wait: config.ack_wait(),
max_deliver: config.max_deliver,
};
let retry_handle = tokio::spawn(async move {
retry_bus
.consume_envelopes::<InboundPayload, _, _>(retry_spec, |_message| async move {
anyhow::bail!("force redelivery")
})
.await
});
bus.publish_envelope(subject.as_str(), &envelope).await?;
let observed = tokio::time::timeout(std::time::Duration::from_secs(20), dlq_rx.recv())
.await?
.ok_or_else(|| anyhow::anyhow!("dlq consumer closed before receiving test message"))?;
retry_handle.abort();
dlq_handle.abort();
assert_eq!(observed.as_deref(), Some(envelope.message_id.as_str()));
Ok(())
}
#[tokio::test]
#[ignore = "requires a running NATS JetStream server from docker-compose.nats.yml"]
async fn jetstream_concurrent_dlq_does_not_lose_messages() -> anyhow::Result<()> {
let mut config = JetStreamConfig::from_env();
config.ack_wait_secs = 1;
config.max_deliver = 2;
config.consumer_concurrency = 4;
let bus = JetStreamBus::connect(config.clone()).await?;
bus.ensure_streams().await?;
let names = StreamNames::new(config.stream_prefix.clone());
let batch = uuid::Uuid::new_v4().to_string();
let count = 8usize;
let account_id = format!("cdlq-{batch}");
// 发布 count 条不同 message_id 的消息
let mut message_ids = Vec::new();
for index in 0..count {
let external = format!("{batch}-{index}");
let subject = inbound_subject("wechat", &account_id);
let payload = InboundPayload {
channel: "wechat".to_string(),
account_id: account_id.clone(),
from_user_id: "tester".to_string(),
content: "force retry".to_string(),
context: None,
external_message_id: Some(external.clone()),
session_id: None,
group_id: None,
create_time_ms: None,
};
let message_id = format!("wechat/{account_id}/{external}");
message_ids.push(message_id.clone());
let envelope = MessageEnvelope::new(
subject.as_str(),
message_id,
format!("trace-{batch}-{index}"),
payload,
);
bus.publish_envelope(subject.as_str(), &envelope).await?;
}
// DLQ consumer 收集 original_message_id
let (dlq_tx, mut dlq_rx) = tokio::sync::mpsc::channel(count);
let dlq_bus = bus.clone();
let dlq_stream = names.name(StreamKind::Dlq);
let dlq_ack_wait = config.ack_wait();
let dlq_max_deliver = config.max_deliver;
let dlq_durable = format!("integration-cdlq-{batch}");
let dlq_handle = tokio::spawn(async move {
dlq_bus
.consume_envelopes::<DeadLetterPayload, _, _>(
ConsumerSpec {
stream: dlq_stream,
durable_name: dlq_durable,
filter_subject: "dlq.>".to_string(),
ack_wait: dlq_ack_wait,
max_deliver: dlq_max_deliver,
},
move |message| {
let dlq_tx = dlq_tx.clone();
async move {
let _ = dlq_tx.send(message.payload.original_message_id).await;
Ok(())
}
},
)
.await
});
// 失败 consumer(总是失败,触发 redelivery -> DLQ
let retry_bus = bus.clone();
let retry_spec = ConsumerSpec {
stream: names.name(StreamKind::Inbound),
durable_name: format!("integration-cdlq-retry-{batch}"),
filter_subject: format!("inbound.wechat.{account_id}"),
ack_wait: config.ack_wait(),
max_deliver: config.max_deliver,
};
let retry_handle = tokio::spawn(async move {
retry_bus
.consume_envelopes::<InboundPayload, _, _>(retry_spec, |_| async move {
anyhow::bail!("force redelivery")
})
.await
});
// 收集 count 条 DLQ 消息,验证无碰撞丢失
let mut received = Vec::new();
for _ in 0..count {
let id = tokio::time::timeout(std::time::Duration::from_secs(30), dlq_rx.recv())
.await?
.ok_or_else(|| anyhow::anyhow!("dlq consumer closed early"))?;
if let Some(id) = id {
received.push(id);
}
}
retry_handle.abort();
dlq_handle.abort();
received.sort();
message_ids.sort();
assert_eq!(received.len(), count, "DLQ 丢失了消息");
for id in &message_ids {
assert!(received.contains(id), "DLQ 缺少消息 {id}");
}
Ok(())
}