Skip to main content

av_bridge/
nats_bus.rs

1//! NATS JetStream connector (brief Module F: "lighter-weight edge/embedded
2//! alternative"). Feature `nats`. One JetStream stream per topic
3//! (`av_<topic>` with dots mapped to underscores), one subject per partition
4//! (`<topic>.p<N>`), giving the same partitioned-ordered-replay contract as
5//! the embedded broker.
6//!
7//! Contract tests gate on `AV_NATS_URL` and skip loudly when unset.
8
9use crate::bus::{partition_for, BusError, EventBus, PublishAck, StoredEvent};
10use crate::manifest::BridgeManifest;
11use std::collections::HashMap;
12
13/// NATS JetStream bus.
14pub 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
26/// Pair up broker credentials, refusing half-configured auth loudly. A
27/// typo'd or unexported password must not silently downgrade the
28/// connection to anonymous (D13: a dropped credential is a silent error;
29/// same contract as the Kafka connector's `KafkaSecurity`).
30fn 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    /// Connect and provision streams per the manifest.
45    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        // rustls 0.23 resolves its process-level CryptoProvider lazily and
53        // panics on first TLS use when more than one provider feature is
54        // compiled in (all-feature workspace builds carry both `ring` and
55        // `aws-lc-rs`) and none has been installed. Pin `ring` here; if the
56        // embedding application already installed one, keep theirs.
57        let _ = rustls_tls::crypto::ring::default_provider().install_default();
58        let url = url.to_owned();
59        // TLS/auth material comes from the environment so `tls://` endpoints
60        // work against private-CA deployments without new constructor
61        // surface: `AV_NATS_CA_FILE` pins a root CA (self-hosted enclaves
62        // rarely use WebPKI certs), `AV_NATS_USER`/`AV_NATS_PASSWORD` supply
63        // broker auth. Values are paths/identifiers, never logged.
64        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                // A pinned CA or supplied credentials both state intent:
73                // this connection must be TLS. Forcing the requirement means
74                // a `nats://` (instead of `tls://`) endpoint typo cannot
75                // silently downgrade to plaintext, and an active MITM
76                // cannot strip `tls_required` from INFO to capture the
77                // CONNECT password in cleartext. A CA file is deliberately
78                // not required for credentials — WebPKI-certified `tls://`
79                // endpoints are legitimate.
80                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        // `async_nats::jetstream::new` must be called from within a Tokio
94        // runtime — the constructor eagerly calls `Handle::current()` and
95        // panics otherwise (async-nats 0.50, jetstream/context.rs:129).
96        // Route it through the executor so the sync `provision` entrypoint
97        // stays runtime-agnostic for callers.
98        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            // Round-21 F8: match the KafkaBus retention arithmetic
107            // discipline. Today `hot_hours: u32` × 3600 fits in
108            // u64, but a future field-widening (e.g., u64 for
109            // very-long-retention research clusters) would
110            // silently wrap here and surface as `Overflow` on the
111            // Kafka path — the same cross-backend divergence
112            // round-20 F1 closed for counters. Use checked_mul
113            // now so a future widening surfaces the error
114            // consistently.
115            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                        // Ephemeral consumers without an inactive_threshold
294                        // linger on the JetStream server until the client
295                        // disconnects. Since our async_nats client stays
296                        // connected for the lifetime of the process, every
297                        // fetch call otherwise leaks a consumer.
298                        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                    // Surface the error instead of silently continuing on:
312                    // an audit-trail consumer that sees a shorter list than
313                    // expected with no error is an evidence gap. The bad
314                    // record stays on JetStream as forensic evidence.
315                    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}