feat: 拆分 worker 架构并增强 assistant 安全保护
- 新增 daemon/channel/assistant/function/scheduler/logging worker 架构和 JetStream 消息协议 - 修复邮件通知误入 assistant、迁移 checksum、日志 subject 和工具调用无限循环风险 - 新增 assistant 工具轮次熔断、工具结果截断和提示词软收敛规则 - 补充邮件推送关键字、可信发件人配置、升级日志和验证测试
This commit is contained in:
+13
-1
@@ -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"] }
|
||||
|
||||
@@ -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 },
|
||||
}
|
||||
@@ -1,3 +1,5 @@
|
||||
pub mod control;
|
||||
pub mod error;
|
||||
pub mod model;
|
||||
pub mod mq;
|
||||
pub mod queue;
|
||||
|
||||
@@ -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()))
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
};
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
@@ -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>),
|
||||
}
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
Reference in New Issue
Block a user