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

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()
     }
@@ -154,8 +148,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,22 @@
-use super::Bridge;
+use super::{Bridge, js_request};
 use crate::error::map_nats_error;
-use crate::nats::{self, RequestClient, session};
+use crate::nats::{self, FlushClient, PublishClient, 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 + PublishClient + FlushClient,
+    C: GetElapsed,
+    J: JetStreamPublisher + JetStreamConsumerFactory,
+>(
     bridge: &Bridge<N, C, J>,
     args: CloseSessionRequest,
 ) -> Result<CloseSessionResponse> {
@@ -28,17 +33,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,22 @@
-use super::Bridge;
+use super::{Bridge, 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 +33,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());