feat(acp-nats, acp-nats-agent): wire JetStream into bridge and agent

Wire Layer 3 of JetStream integration — the actual message flow.

Bridge prompt handler branches on bridge.js(): when JetStream is
available, publishes to COMMANDS stream and consumes from
NOTIFICATIONS/RESPONSES streams instead of subscribe-before-publish.

Session.ready skips the 100ms sleep when JetStream is available —
the RESPONSES stream captures the message regardless of timing.

Agent library gains serve_js() alongside serve() — creates a consumer
on the COMMANDS stream and dispatches with ack/nak/term signals. Runs
in parallel with core NATS serve via with_jetstream() constructor.

Binary crates (stdio, ws) create NatsJetStreamClient from the NATS
connection and pass to Bridge::with_jetstream().

() implements JetStream traits as no-ops for backward compatibility
when JetStream is not configured.

Signed-off-by: Yordis Prieto <yordis.prieto@gmail.com>

Author: yordis

rsworkspace/crates/acp-nats-agent/src/connection.rs

@@ -20,12 +20,14 @@ use trogon_nats::{FlushClient, PublishClient, RequestClient, SubscribeClient};
 
 pub enum ConnectionError {
     Subscribe(Box<dyn std::error::Error + Send + Sync>),
+    JetStream(Box<dyn std::error::Error + Send + Sync>),
 }
 
 impl std::fmt::Debug for ConnectionError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             Self::Subscribe(e) => f.debug_tuple("Subscribe").field(e).finish(),
+            Self::JetStream(e) => f.debug_tuple("JetStream").field(e).finish(),
         }
     }
 }
@@ -34,6 +36,7 @@ impl std::fmt::Display for ConnectionError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             Self::Subscribe(e) => write!(f, "failed to subscribe: {}", e),
+            Self::JetStream(e) => write!(f, "jetstream error: {}", e),
         }
     }
 }
@@ -95,6 +98,47 @@ where
         (conn, io_task)
     }
 
+    pub fn with_jetstream<J>(
+        agent: impl Agent + 'static,
+        nats: N,
+        js: J,
+        acp_prefix: AcpPrefix,
+        spawn: impl Fn(LocalBoxFuture<'static, ()>) + Copy + 'static,
+    ) -> (
+        Self,
+        impl std::future::Future<Output = Result<(), ConnectionError>>,
+    )
+    where
+        J: JetStreamConsumerFactory + 'static,
+    {
+        let nats_for_serve = nats.clone();
+        let nats_for_js = nats.clone();
+        let prefix = acp_prefix.as_str().to_string();
+        let prefix_js = prefix.clone();
+
+        let io_task = async move {
+            let (agent1, agent2) = {
+                let agent = Rc::new(agent);
+                (agent.clone(), agent)
+            };
+
+            let core = serve(agent1, nats_for_serve, &prefix, spawn);
+            let jetstream = serve_js(agent2, nats_for_js, js, &prefix_js, spawn);
+
+            tokio::select! {
+                result = core => result,
+                result = jetstream => result,
+            }
+        };
+
+        let conn = Self {
+            nats,
+            acp_prefix,
+            operation_timeout: DEFAULT_OPERATION_TIMEOUT,
+        };
+        (conn, io_task)
+    }
+
     pub fn client_for_session(&self, session_id: AcpSessionId) -> NatsClientProxy<N> {
         NatsClientProxy::new(
             self.nats.clone(),
@@ -313,6 +357,183 @@ where
         .map_err(DispatchError::NotificationHandler)
 }
 
+use trogon_nats::jetstream::{JetStreamConsumer as _, JetStreamConsumerFactory, JsMessage};
+
+async fn serve_js<N, J, A>(
+    agent: A,
+    nats: N,
+    js: J,
+    prefix: &str,
+    spawn: impl Fn(LocalBoxFuture<'static, ()>) + 'static,
+) -> Result<(), ConnectionError>
+where
+    N: PublishClient + FlushClient + Clone + 'static,
+    J: JetStreamConsumerFactory + 'static,
+    A: Agent + 'static,
+{
+    let stream_name = acp_nats::jetstream::streams::commands_stream_name(prefix);
+    let config = acp_nats::jetstream::consumers::commands_observer();
+
+    info!(stream = %stream_name, "Starting JetStream consumer for COMMANDS stream");
+
+    let consumer = js
+        .create_consumer(&stream_name, config)
+        .await
+        .map_err(|e| ConnectionError::JetStream(Box::new(e)))?;
+
+    let mut messages = consumer
+        .messages()
+        .await
+        .map_err(|e| ConnectionError::JetStream(Box::new(e)))?;
+
+    let agent = Rc::new(agent);
+    let nats = Rc::new(nats);
+
+    while let Some(msg_result) = messages.next().await {
+        match msg_result {
+            Ok(js_msg) => {
+                let agent = agent.clone();
+                let nats = nats.clone();
+                spawn(Box::pin(async move {
+                    dispatch_js_message(js_msg, agent.as_ref(), nats.as_ref()).await;
+                }));
+            }
+            Err(e) => {
+                warn!(error = %e, "JetStream consumer error");
+            }
+        }
+    }
+
+    info!("JetStream COMMANDS consumer ended");
+    Ok(())
+}
+
+async fn dispatch_js_message<N: PublishClient + FlushClient, A: Agent>(
+    js_msg: JsMessage,
+    agent: &A,
+    nats: &N,
+) {
+    let msg = js_msg.message();
+    let subject = msg.subject.as_str();
+
+    let parsed = match parse_agent_subject(subject) {
+        Some(p) => p,
+        None => {
+            if let Err(e) = js_msg.term().await {
+                warn!(error = %e, subject, "Failed to term unknown subject");
+            }
+            return;
+        }
+    };
+
+    let result = match parsed.method {
+        AgentMethod::Initialize => {
+            handle_request(msg, nats, |req: InitializeRequest| agent.initialize(req)).await
+        }
+        AgentMethod::Authenticate => {
+            handle_request(msg, nats, |req: AuthenticateRequest| {
+                agent.authenticate(req)
+            })
+            .await
+        }
+        AgentMethod::SessionNew => {
+            handle_request(msg, nats, |req: NewSessionRequest| agent.new_session(req)).await
+        }
+        AgentMethod::SessionList => {
+            handle_request(msg, nats, |req: ListSessionsRequest| {
+                agent.list_sessions(req)
+            })
+            .await
+        }
+        AgentMethod::SessionLoad => {
+            handle_request(msg, nats, |req: LoadSessionRequest| agent.load_session(req)).await
+        }
+        AgentMethod::SessionPrompt => {
+            handle_request(msg, nats, |req: PromptRequest| agent.prompt(req)).await
+        }
+        AgentMethod::SessionCancel => {
+            handle_notification(msg, |req: CancelNotification| agent.cancel(req)).await
+        }
+        AgentMethod::SessionSetMode => {
+            handle_request(msg, nats, |req: SetSessionModeRequest| {
+                agent.set_session_mode(req)
+            })
+            .await
+        }
+        AgentMethod::SessionSetConfigOption => {
+            handle_request(msg, nats, |req: SetSessionConfigOptionRequest| {
+                agent.set_session_config_option(req)
+            })
+            .await
+        }
+        AgentMethod::SessionSetModel => {
+            handle_request(msg, nats, |req: SetSessionModelRequest| {
+                agent.set_session_model(req)
+            })
+            .await
+        }
+        AgentMethod::SessionFork => {
+            handle_request(msg, nats, |req: ForkSessionRequest| agent.fork_session(req)).await
+        }
+        AgentMethod::SessionResume => {
+            handle_request(msg, nats, |req: ResumeSessionRequest| {
+                agent.resume_session(req)
+            })
+            .await
+        }
+        AgentMethod::SessionClose => {
+            handle_request(msg, nats, |req: CloseSessionRequest| {
+                agent.close_session(req)
+            })
+            .await
+        }
+        AgentMethod::Ext(_) => {
+            if msg.reply.is_some() {
+                handle_request(msg, nats, |req: ExtRequest| agent.ext_method(req)).await
+            } else {
+                handle_notification(msg, |req: ExtNotification| agent.ext_notification(req)).await
+            }
+        }
+    };
+
+    match &result {
+        Ok(()) => {
+            if let Err(e) = js_msg.ack().await {
+                warn!(subject, error = %e, "Failed to ack JetStream message");
+            }
+        }
+        Err(DispatchError::DeserializeRequest(_) | DispatchError::DeserializeNotification(_)) => {
+            if let Err(e) = js_msg.term().await {
+                warn!(subject, error = %e, "Failed to term bad payload");
+            }
+        }
+        Err(DispatchError::NoReplySubject) => {
+            if let Err(e) = js_msg.term().await {
+                warn!(subject, error = %e, "Failed to term missing reply subject");
+            }
+        }
+        Err(DispatchError::Reply(_)) => {
+            if let Err(e) = js_msg.nak().await {
+                warn!(subject, error = %e, "Failed to nak after reply failure");
+            }
+        }
+        Err(DispatchError::NotificationHandler(_)) => {
+            if let Err(e) = js_msg.ack().await {
+                warn!(subject, error = %e, "Failed to ack after notification handler error");
+            }
+        }
+    }
+
+    if let Err(e) = result {
+        let sid = parsed
+            .session_id
+            .as_ref()
+            .map(|s| s.as_str())
+            .unwrap_or("-");
+        warn!(subject, session_id = sid, error = %e, "Error handling agent request");
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

rsworkspace/crates/acp-nats-stdio/src/main.rs

@@ -32,10 +32,14 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let stdin = async_compat::Compat::new(tokio::io::stdin());
     let stdout = async_compat::Compat::new(tokio::io::stdout());
 
+    let js_context = async_nats::jetstream::new(nats_client.clone());
+    let js = acp_nats::NatsJetStreamClient::new(js_context);
+
     let local = tokio::task::LocalSet::new();
     let result = local
         .run_until(run_bridge(
             nats_client,
+            js,
             &config,
             stdout,
             stdin,
@@ -57,8 +61,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 #[cfg(coverage)]
 fn main() {}
 
-async fn run_bridge<N, W, R>(
+async fn run_bridge<N, J, W, R>(
     nats_client: N,
+    js: J,
     config: &acp_nats::Config,
     stdout: W,
     stdin: R,
@@ -70,13 +75,16 @@ where
         + acp_nats::FlushClient
         + acp_nats::SubscribeClient
         + 'static,
+    J: acp_nats::JetStreamPublisher + acp_nats::JetStreamConsumerFactory + 'static,
     W: futures::AsyncWrite + Unpin + 'static,
     R: futures::AsyncRead + Unpin + 'static,
 {
     let meter = acp_telemetry::meter("acp-io-bridge-nats");
     let (notification_tx, notification_rx) = tokio::sync::mpsc::channel::<SessionNotification>(64);
-    let bridge = Rc::new(Bridge::new(
+
+    let bridge = Rc::new(Bridge::with_jetstream(
         nats_client.clone(),
+        js,
         SystemClock,
         &meter,
         config.clone(),
@@ -166,6 +174,7 @@ mod tests {
         let result = local
             .run_until(run_bridge(
                 mock,
+                (),
                 &config,
                 stdout,
                 stdin,
@@ -198,6 +207,7 @@ mod tests {
         let result = local
             .run_until(run_bridge(
                 mock,
+                (),
                 &config,
                 stdout,
                 stdin,

rsworkspace/crates/acp-nats-ws/src/connection.rs

@@ -1,3 +1,4 @@
+use acp_nats::{JetStreamConsumerFactory, JetStreamPublisher};
 use acp_nats::{StdJsonSerialize, agent::Bridge, client, spawn_notification_forwarder};
 use agent_client_protocol::{AgentSideConnection, SessionNotification};
 use axum::extract::ws::{Message, WebSocket};
@@ -12,9 +13,10 @@ use trogon_std::time::SystemClock;
 use crate::constants::DUPLEX_BUFFER_SIZE;
 
 /// Handles a single WebSocket connection by bridging it to NATS via ACP.
-pub async fn handle<N>(
+pub async fn handle<N, J>(
     socket: WebSocket,
     nats_client: N,
+    js: J,
     config: acp_nats::Config,
     mut shutdown_rx: watch::Receiver<bool>,
 ) where
@@ -24,6 +26,7 @@ pub async fn handle<N>(
         + acp_nats::SubscribeClient
         + Clone
         + 'static,
+    J: JetStreamPublisher + JetStreamConsumerFactory + 'static,
 {
     let (ws_sender, ws_receiver) = socket.split();
 
@@ -35,8 +38,9 @@ pub async fn handle<N>(
 
     let meter = acp_telemetry::meter("acp-nats-ws");
     let (notification_tx, notification_rx) = tokio::sync::mpsc::channel::<SessionNotification>(64);
-    let bridge = Rc::new(Bridge::new(
+    let bridge = Rc::new(Bridge::with_jetstream(
         nats_client.clone(),
+        js,
         SystemClock,
         &meter,
         config,