1use crate::bus::{partition_for, BusError, EventBus, PublishAck, StoredEvent};
10use crate::manifest::BridgeManifest;
11use std::collections::HashMap;
12
13pub struct NatsBus {
15 cold_archive: Option<crate::cold_store::ColdArchive>,
16 js: async_nats::jetstream::Context,
17 executor: crate::bus::ConnectorExecutor,
18 topics: HashMap<String, u32>,
19 validators: HashMap<String, jsonschema::Validator>,
20}
21
22fn stream_name(topic: &str) -> String {
23 format!("av_{}", topic.replace('.', "_"))
24}
25
26fn nats_credentials(
31 user: Option<String>,
32 password: Option<String>,
33) -> Result<Option<(String, String)>, BusError> {
34 match (user, password) {
35 (Some(user), Some(password)) => Ok(Some((user, password))),
36 (None, None) => Ok(None),
37 _ => Err(BusError::Backend(
38 "AV_NATS_USER and AV_NATS_PASSWORD must be set together".to_owned(),
39 )),
40 }
41}
42
43impl NatsBus {
44 pub fn provision(url: &str, manifest: &BridgeManifest) -> Result<Self, BusError> {
46 manifest
47 .validate()
48 .map_err(|e| BusError::Backend(e.to_string()))?;
49 let validators = crate::manifest::compile_topic_validators(manifest)?;
50 let cold_archive = crate::cold_store::ColdArchive::from_manifest(manifest)?;
51 let executor = crate::bus::ConnectorExecutor::new("agentvisor-ai-nats")?;
52 let _ = rustls_tls::crypto::ring::default_provider().install_default();
58 let url = url.to_owned();
59 let ca_file = std::env::var_os("AV_NATS_CA_FILE").map(std::path::PathBuf::from);
65 let credentials = nats_credentials(
66 std::env::var("AV_NATS_USER").ok(),
67 std::env::var("AV_NATS_PASSWORD").ok(),
68 )?;
69 let client = executor
70 .run(move || async move {
71 let mut options = async_nats::ConnectOptions::new();
72 let secured = ca_file.is_some() || credentials.is_some();
81 if let Some(ca) = ca_file {
82 options = options.add_root_certificates(ca);
83 }
84 if secured {
85 options = options.require_tls(true);
86 }
87 if let Some((user, password)) = credentials {
88 options = options.user_and_password(user, password);
89 }
90 options.connect(url).await
91 })?
92 .map_err(|e| BusError::Backend(e.to_string()))?;
93 let js = executor.run({
99 let client = client.clone();
100 move || async move { async_nats::jetstream::new(client) }
101 })?;
102 let mut topics = HashMap::new();
103 for t in &manifest.topics {
104 let subjects: Vec<String> = (0..t.partitions).map(|p| format!("{}.p{p}", t.name)).collect();
105 let context = js.clone();
106 let retention_secs = u64::from(t.retention.hot_hours)
116 .checked_mul(3600)
117 .ok_or_else(|| BusError::Backend("NATS retention overflow".to_owned()))?;
118 let retention = std::time::Duration::from_secs(retention_secs);
119 let config = async_nats::jetstream::stream::Config {
120 name: stream_name(&t.name),
121 subjects: subjects.clone(),
122 max_age: retention,
123 duplicate_window: retention,
124 num_replicas: usize::try_from(manifest.replication_factor)
125 .map_err(|_| BusError::Backend("NATS replication factor exceeds usize".to_owned()))?,
126 ..Default::default()
127 };
128 executor
129 .run(move || async move {
130 context
131 .get_or_create_stream(config.clone())
132 .await
133 .map_err(|error| error.to_string())?;
134 context
135 .update_stream(config.clone())
136 .await
137 .map_err(|error| error.to_string())?;
138 let mut stream = context
139 .get_stream(&config.name)
140 .await
141 .map_err(|error| error.to_string())?;
142 let actual = stream
143 .info()
144 .await
145 .map_err(|error| error.to_string())?
146 .config
147 .clone();
148 if actual.subjects != subjects
149 || actual.max_age != retention
150 || actual.duplicate_window != retention
151 || actual.num_replicas != config.num_replicas
152 {
153 return Err("JetStream stream does not match Bridge manifest".to_owned());
154 }
155 Ok::<_, String>(())
156 })?
157 .map_err(BusError::Backend)?;
158 topics.insert(t.name.clone(), t.partitions);
159 }
160 Ok(Self {
161 cold_archive,
162 js,
163 executor,
164 topics,
165 validators,
166 })
167 }
168
169 fn publish_with_uid(
170 &self,
171 topic: &str,
172 key: &str,
173 value: &serde_json::Value,
174 event_uid: Option<&str>,
175 ) -> Result<PublishAck, BusError> {
176 crate::manifest::validate_topic_event(&self.validators, topic, value)?;
177 let partitions = *self
178 .topics
179 .get(topic)
180 .ok_or_else(|| BusError::UnknownTopic(topic.to_owned()))?;
181 let partition = partition_for(key, partitions);
182 let event_uid = event_uid.map_or_else(av_core::new_event_uid, str::to_owned);
183 let stored_at = av_core::time::now_ms();
184 let record = StoredEvent {
185 partition,
186 offset: 0,
187 key: key.to_owned(),
188 value: value.clone(),
189 stored_at,
190 };
191 if let Some(archive) = &self.cold_archive {
192 archive.stage(topic, &record, &event_uid)?;
193 }
194 let ack = self.publish_broker_only(topic, key, value, stored_at, &event_uid)?;
195 if let Some(archive) = &self.cold_archive {
196 archive.commit(topic, &event_uid, ack.offset)?;
197 }
198 Ok(ack)
199 }
200
201 fn publish_broker_only(
202 &self,
203 topic: &str,
204 key: &str,
205 value: &serde_json::Value,
206 stored_at: u64,
207 event_uid: &str,
208 ) -> Result<PublishAck, BusError> {
209 crate::manifest::validate_topic_event(&self.validators, topic, value)?;
210 let partitions = *self
211 .topics
212 .get(topic)
213 .ok_or_else(|| BusError::UnknownTopic(topic.to_owned()))?;
214 let partition = partition_for(key, partitions);
215 let record = StoredEvent {
216 partition,
217 offset: 0,
218 key: key.to_owned(),
219 value: value.clone(),
220 stored_at,
221 };
222 let payload = serde_json::to_vec(&record)?;
223 let subject = format!("{topic}.p{partition}");
224 let context = self.js.clone();
225 let event_uid = event_uid.to_owned();
226 let ack = self
227 .executor
228 .run(move || async move {
229 let mut headers = async_nats::HeaderMap::new();
230 headers.insert("Nats-Msg-Id", event_uid);
231 context
232 .publish_with_headers(subject, headers, payload.into())
233 .await?
234 .await
235 })?
236 .map_err(|e| BusError::Backend(e.to_string()))?;
237 Ok(PublishAck {
238 topic: topic.to_owned(),
239 partition,
240 offset: ack.sequence,
241 })
242 }
243}
244
245impl EventBus for NatsBus {
246 fn set_control_key(&self, key: [u8; 32]) -> Result<(), BusError> {
247 if let Some(archive) = &self.cold_archive {
248 archive.set_control_key(key)?;
249 }
250 Ok(())
251 }
252
253 fn publish(&self, topic: &str, key: &str, value: &serde_json::Value) -> Result<PublishAck, BusError> {
254 let event_uid = value
255 .get("metadata")
256 .and_then(|metadata| metadata.get("uid"))
257 .and_then(serde_json::Value::as_str);
258 self.publish_with_uid(topic, key, value, event_uid)
259 }
260
261 fn publish_idempotent(
262 &self,
263 topic: &str,
264 key: &str,
265 value: &serde_json::Value,
266 event_uid: &str,
267 ) -> Result<PublishAck, BusError> {
268 self.publish_with_uid(topic, key, value, Some(event_uid))
269 }
270
271 fn fetch(
272 &self,
273 topic: &str,
274 partition: u32,
275 offset: u64,
276 max: usize,
277 ) -> Result<Vec<StoredEvent>, BusError> {
278 if !self.topics.contains_key(topic) {
279 return Err(BusError::UnknownTopic(topic.to_owned()));
280 }
281 let stream_name = stream_name(topic);
282 let subject = format!("{topic}.p{partition}");
283 let context = self.js.clone();
284 self.executor
285 .run(move || async move {
286 let stream = context.get_stream(&stream_name).await?;
287 let consumer = stream
288 .create_consumer(async_nats::jetstream::consumer::pull::Config {
289 deliver_policy: async_nats::jetstream::consumer::DeliverPolicy::ByStartSequence {
290 start_sequence: offset.max(1),
291 },
292 filter_subject: subject,
293 inactive_threshold: std::time::Duration::from_secs(30),
299 ..Default::default()
300 })
301 .await?;
302 let mut out = Vec::new();
303 let mut batch = consumer
304 .batch()
305 .max_messages(max)
306 .expires(std::time::Duration::from_millis(500))
307 .messages()
308 .await?;
309 use futures::StreamExt as _;
310 while let Some(next) = batch.next().await {
311 let msg = next?;
316 let mut ev: StoredEvent = serde_json::from_slice(&msg.payload).map_err(|error| {
317 async_nats::Error::from(std::io::Error::other(format!("fetch decode: {error}")))
318 })?;
319 if let Ok(info) = msg.info() {
320 ev.offset = info.stream_sequence;
321 }
322 out.push(ev);
323 if out.len() >= max {
324 break;
325 }
326 }
327 Ok::<_, async_nats::Error>(out)
328 })?
329 .map_err(|e| BusError::Backend(e.to_string()))
330 }
331
332 fn partitions(&self, topic: &str) -> Result<u32, BusError> {
333 self.topics
334 .get(topic)
335 .copied()
336 .ok_or_else(|| BusError::UnknownTopic(topic.to_owned()))
337 }
338
339 fn topics(&self) -> Vec<String> {
340 let mut t: Vec<String> = self.topics.keys().cloned().collect();
341 t.sort();
342 t
343 }
344
345 fn maintenance(&self, _now_ms: u64) -> Result<u64, BusError> {
346 self.cold_archive.as_ref().map_or(Ok(0), |archive| {
347 archive.retry_pending_with(|pending| {
348 self.publish_broker_only(
349 &pending.topic,
350 &pending.key,
351 &pending.value,
352 pending.stored_at,
353 &pending.event_uid,
354 )
355 })
356 })
357 }
358}
359
360#[cfg(test)]
361mod tests {
362 #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
363
364 use super::nats_credentials;
365
366 #[test]
367 fn credentials_pair_or_fail_loudly() {
368 assert!(nats_credentials(None, None).unwrap().is_none());
369 assert_eq!(
370 nats_credentials(Some("u".into()), Some("p".into())).unwrap(),
371 Some(("u".into(), "p".into()))
372 );
373 for (user, password) in [(Some("u".to_owned()), None), (None, Some("p".to_owned()))] {
374 let error = nats_credentials(user, password).unwrap_err();
375 assert!(
376 error.to_string().contains("must be set together"),
377 "partial credentials must fail loudly, got: {error}"
378 );
379 }
380 }
381}