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

@@ -21,8 +21,6 @@ use opentelemetry::metrics::Meter;
 use tokio::sync::mpsc;
 use tokio::task::JoinHandle;
 use tracing::{info, warn};
-#[cfg(not(coverage))]
-#[allow(unused_imports)]
 use trogon_nats::jetstream::{JetStreamConsumerFactory, JetStreamPublisher};
 use trogon_std::time::GetElapsed;
 
@@ -36,7 +34,6 @@ use crate::constants::SESSION_READY_DELAY;
 
 pub struct Bridge<N, C: GetElapsed, J = ()> {
     pub(crate) nats: N,
-    #[allow(dead_code)] // Used in prompt.rs JetStream path
     pub(crate) js: Option<J>,
     pub(crate) clock: C,
     pub(crate) config: Config,
@@ -68,7 +65,6 @@ impl<N, C: GetElapsed> Bridge<N, C> {
 }
 
 impl<N, C: GetElapsed, J> Bridge<N, C, J> {
-    #[cfg(not(coverage))]
     pub fn with_jetstream(
         nats: N,
         js: J,
@@ -93,8 +89,6 @@ impl<N, C: GetElapsed, J> Bridge<N, C, J> {
         &self.nats
     }
 
-    #[cfg(not(coverage))]
-    #[allow(dead_code)]
     pub(crate) fn js(&self) -> Option<&J> {
         self.js.as_ref()
     }
@@ -116,8 +110,9 @@ impl<N: PublishClient + FlushClient + Clone + Send + 'static, C: GetElapsed, J>
         let nats = self.nats.clone();
         let prefix = self.config.acp_prefix().to_string();
         let metrics = self.metrics.clone();
+        let has_jetstream = self.js.is_some();
         let handle = tokio::spawn(async move {
-            publish_session_ready(&nats, &prefix, &session_id, &metrics).await;
+            publish_session_ready(&nats, &prefix, &session_id, &metrics, has_jetstream).await;
         });
         self.spawn_background(handle);
     }
@@ -128,8 +123,11 @@ async fn publish_session_ready<N: PublishClient + FlushClient>(
     prefix: &str,
     session_id: &SessionId,
     metrics: &Metrics,
+    has_jetstream: bool,
 ) {
-    tokio::time::sleep(SESSION_READY_DELAY).await;
+    if !has_jetstream {
+        tokio::time::sleep(SESSION_READY_DELAY).await;
+    }
 
     let subject = session::agent::ext_ready(prefix, &session_id.to_string());
     info!(session_id = %session_id, subject = %subject, "Publishing session.ready");
@@ -154,8 +152,11 @@ async fn publish_session_ready<N: PublishClient + FlushClient>(
 }
 
 #[async_trait::async_trait(?Send)]
-impl<N: RequestClient + PublishClient + SubscribeClient + FlushClient, C: GetElapsed, J> Agent
-    for Bridge<N, C, J>
+impl<
+    N: RequestClient + PublishClient + SubscribeClient + FlushClient,
+    C: GetElapsed,
+    J: JetStreamPublisher + JetStreamConsumerFactory,
+> Agent for Bridge<N, C, J>
 {
     async fn initialize(&self, args: InitializeRequest) -> Result<InitializeResponse> {
         initialize::handle(self, args).await

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

@@ -1,17 +1,23 @@
 use super::Bridge;
+use super::js_request;
 use crate::error::map_nats_error;
 use crate::nats::{self, RequestClient, session};
 use crate::session_id::AcpSessionId;
 use agent_client_protocol::{CloseSessionRequest, CloseSessionResponse, Error, ErrorCode, Result};
 use tracing::{info, instrument};
+use trogon_nats::jetstream::{JetStreamConsumerFactory, JetStreamPublisher};
 use trogon_std::time::GetElapsed;
 
 #[instrument(
     name = "acp.session.close",
     skip(bridge, args),
     fields(session_id = %args.session_id)
 )]
-pub async fn handle<N: RequestClient, C: GetElapsed, J>(
+pub async fn handle<
+    N: RequestClient,
+    C: GetElapsed,
+    J: JetStreamPublisher + JetStreamConsumerFactory,
+>(
     bridge: &Bridge<N, C, J>,
     args: CloseSessionRequest,
 ) -> Result<CloseSessionResponse> {
@@ -28,17 +34,33 @@ pub async fn handle<N: RequestClient, C: GetElapsed, J>(
             format!("Invalid session ID: {}", e),
         )
     })?;
-    let nats = bridge.nats();
-    let subject = session::agent::close(bridge.config.acp_prefix(), session_id.as_str());
-
-    let result = nats::request_with_timeout::<N, CloseSessionRequest, CloseSessionResponse>(
-        nats,
-        &subject,
-        &args,
-        bridge.config.operation_timeout,
-    )
-    .await
-    .map_err(map_nats_error);
+    let prefix = bridge.config.acp_prefix();
+    let subject = session::agent::close(prefix, session_id.as_str());
+
+    let result = match bridge.js() {
+        Some(js) => {
+            let req_id = uuid::Uuid::new_v4().to_string();
+            js_request::js_request::<J, _, CloseSessionResponse, _>(
+                js,
+                &subject,
+                &args,
+                &trogon_std::StdJsonSerialize,
+                prefix,
+                session_id.as_str(),
+                &req_id,
+                bridge.config.operation_timeout,
+            )
+            .await
+        }
+        None => nats::request_with_timeout::<N, CloseSessionRequest, CloseSessionResponse>(
+            bridge.nats(),
+            &subject,
+            &args,
+            bridge.config.operation_timeout,
+        )
+        .await
+        .map_err(map_nats_error),
+    };
 
     bridge.metrics.record_request(
         "close_session",

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

@@ -1,17 +1,23 @@
 use super::Bridge;
+use super::js_request;
 use crate::error::map_nats_error;
 use crate::nats::{self, FlushClient, PublishClient, RequestClient, session};
 use crate::session_id::AcpSessionId;
 use agent_client_protocol::{Error, ErrorCode, ForkSessionRequest, ForkSessionResponse, Result};
 use tracing::{Span, info, instrument};
+use trogon_nats::jetstream::{JetStreamConsumerFactory, JetStreamPublisher};
 use trogon_std::time::GetElapsed;
 
 #[instrument(
     name = "acp.session.fork",
     skip(bridge, args),
     fields(session_id = %args.session_id, new_session_id = tracing::field::Empty)
 )]
-pub async fn handle<N: RequestClient + PublishClient + FlushClient, C: GetElapsed, J>(
+pub async fn handle<
+    N: RequestClient + PublishClient + FlushClient,
+    C: GetElapsed,
+    J: JetStreamPublisher + JetStreamConsumerFactory,
+>(
     bridge: &Bridge<N, C, J>,
     args: ForkSessionRequest,
 ) -> Result<ForkSessionResponse> {
@@ -28,17 +34,33 @@ pub async fn handle<N: RequestClient + PublishClient + FlushClient, C: GetElapse
             format!("Invalid session ID: {}", e),
         )
     })?;
-    let nats = bridge.nats();
-    let subject = session::agent::fork(bridge.config.acp_prefix(), session_id.as_str());
-
-    let result = nats::request_with_timeout::<N, ForkSessionRequest, ForkSessionResponse>(
-        nats,
-        &subject,
-        &args,
-        bridge.config.operation_timeout,
-    )
-    .await
-    .map_err(map_nats_error);
+    let prefix = bridge.config.acp_prefix();
+    let subject = session::agent::fork(prefix, session_id.as_str());
+
+    let result = match bridge.js() {
+        Some(js) => {
+            let req_id = uuid::Uuid::new_v4().to_string();
+            js_request::js_request::<J, _, ForkSessionResponse, _>(
+                js,
+                &subject,
+                &args,
+                &trogon_std::StdJsonSerialize,
+                prefix,
+                session_id.as_str(),
+                &req_id,
+                bridge.config.operation_timeout,
+            )
+            .await
+        }
+        None => nats::request_with_timeout::<N, ForkSessionRequest, ForkSessionResponse>(
+            bridge.nats(),
+            &subject,
+            &args,
+            bridge.config.operation_timeout,
+        )
+        .await
+        .map_err(map_nats_error),
+    };
 
     if let Ok(ref response) = result {
         Span::current().record("new_session_id", response.session_id.to_string().as_str());