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

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

JsMessage is now a zero-cost trait abstraction: NatsJsMessage wraps real
jetstream::Message for production, MockJsMessage records signals for
testing. All dispatch and prompt code is generic over the trait.

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.

Agent library gains serve_js() alongside serve() with full
ack/nak/term signal handling. Runs in parallel via with_jetstream().

Binary crates (stdio, ws) create NatsJetStreamClient and pass to
Bridge::with_jetstream(). () implements JetStream traits for backward
compatibility.

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,253 @@ where
         .map_err(DispatchError::NotificationHandler)
 }
 
+use trogon_nats::jetstream::{JetStreamConsumer as _, JetStreamConsumerFactory, JsMessage};
+
+const KEEPALIVE_INTERVAL: Duration = Duration::from_secs(15);
+
+async fn handle_request_with_keepalive<N, Resp, ReqT, F, M>(
+    msg: &Message,
+    nats: &N,
+    js_msg: &M,
+    handler: impl FnOnce(ReqT) -> F,
+) -> Result<(), DispatchError>
+where
+    N: PublishClient + FlushClient,
+    ReqT: serde::de::DeserializeOwned,
+    F: std::future::Future<Output = agent_client_protocol::Result<Resp>>,
+    Resp: serde::Serialize,
+    M: JsMessage,
+{
+    let reply_to = msg.reply.as_deref().ok_or(DispatchError::NoReplySubject)?;
+
+    let request: ReqT = match serde_json::from_slice(&msg.payload) {
+        Ok(req) => req,
+        Err(e) => {
+            let error = agent_client_protocol::Error::new(
+                agent_client_protocol::ErrorCode::InvalidParams.into(),
+                format!("Failed to deserialize request: {}", e),
+            );
+            let _ = reply(nats, reply_to, &error).await;
+            return Err(DispatchError::DeserializeRequest(e));
+        }
+    };
+
+    let handler_fut = handler(request);
+    tokio::pin!(handler_fut);
+
+    let mut keepalive = tokio::time::interval(KEEPALIVE_INTERVAL);
+    keepalive.tick().await;
+
+    loop {
+        tokio::select! {
+            result = &mut handler_fut => {
+                return match result {
+                    Ok(resp) => reply(nats, reply_to, &resp).await,
+                    Err(err) => reply(nats, reply_to, &err).await,
+                };
+            }
+            _ = keepalive.tick() => {
+                if let Err(e) = js_msg.in_progress().await {
+                    warn!(error = %e, "Failed to send in_progress keepalive");
+                }
+            }
+        }
+    }
+}
+
+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, M: JsMessage>(
+    js_msg: M,
+    agent: &A,
+    nats: &N,
+) {
+    let subject = js_msg.subject().to_string();
+    let msg = Message {
+        subject: subject.as_str().into(),
+        reply: js_msg.reply().map(|s| s.into()),
+        payload: js_msg.payload().clone(),
+        headers: js_msg.headers().cloned(),
+        status: None,
+        description: None,
+        length: js_msg.payload().len(),
+    };
+    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_with_keepalive(&msg, nats, &js_msg, |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,