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

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

@@ -1,3 +1,4 @@
 use std::time::Duration;
 
 pub const DEFAULT_OPERATION_TIMEOUT: Duration = Duration::from_secs(30);
+pub const KEEPALIVE_INTERVAL: Duration = Duration::from_secs(15);

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

@@ -21,22 +21,19 @@ 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_nats::jetstream::{JetStreamConsumerFactory, JetStreamPublisher, JsRequestMessage};
 use trogon_std::time::GetElapsed;
 
 use super::{
     authenticate, cancel, close_session, ext_method, ext_notification, fork_session, initialize,
-    list_sessions, load_session, new_session, prompt, resume_session, set_session_config_option,
-    set_session_mode, set_session_model,
+    js_request, list_sessions, load_session, new_session, prompt, resume_session,
+    set_session_config_option, set_session_mode, set_session_model,
 };
 
 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()
     }
@@ -153,9 +147,61 @@ async fn publish_session_ready<N: PublishClient + FlushClient>(
     }
 }
 
+impl<
+    N: RequestClient + PublishClient + FlushClient,
+    C: GetElapsed,
+    J: JetStreamPublisher + JetStreamConsumerFactory,
+> Bridge<N, C, J>
+where
+    <J::Consumer as trogon_nats::jetstream::JetStreamConsumer>::Message: JsRequestMessage,
+{
+    pub(crate) async fn session_request<Req, Res>(
+        &self,
+        subject: &str,
+        args: &Req,
+        session_id: &str,
+    ) -> Result<Res>
+    where
+        Req: serde::Serialize,
+        Res: serde::de::DeserializeOwned,
+    {
+        use crate::error::map_nats_error;
+
+        match self.js() {
+            Some(js) => {
+                let req_id = uuid::Uuid::new_v4().to_string();
+                js_request::js_request::<J, _, Res, _>(
+                    js,
+                    subject,
+                    args,
+                    &trogon_std::StdJsonSerialize,
+                    self.config.acp_prefix(),
+                    session_id,
+                    &req_id,
+                    self.config.operation_timeout,
+                )
+                .await
+            }
+            None => nats::request_with_timeout::<N, Req, Res>(
+                self.nats(),
+                subject,
+                args,
+                self.config.operation_timeout,
+            )
+            .await
+            .map_err(map_nats_error),
+        }
+    }
+}
+
 #[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>
+where
+    <J::Consumer as trogon_nats::jetstream::JetStreamConsumer>::Message: JsRequestMessage,
 {
     async fn initialize(&self, args: InitializeRequest) -> Result<InitializeResponse> {
         initialize::handle(self, args).await

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

@@ -1,20 +1,27 @@
 use super::Bridge;
-use crate::error::map_nats_error;
-use crate::nats::{self, RequestClient, session};
+use crate::nats::{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, JsRequestMessage};
 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> {
+) -> Result<CloseSessionResponse>
+where
+    <J::Consumer as trogon_nats::jetstream::JetStreamConsumer>::Message: JsRequestMessage,
+{
     let start = bridge.clock.now();
 
     info!(session_id = %args.session_id, "Close session request");
@@ -28,17 +35,16 @@ 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 = bridge
+        .session_request::<CloseSessionRequest, CloseSessionResponse>(
+            &subject,
+            &args,
+            session_id.as_str(),
+        )
+        .await;
 
     bridge.metrics.record_request(
         "close_session",

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

@@ -1,20 +1,27 @@
 use super::Bridge;
-use crate::error::map_nats_error;
-use crate::nats::{self, FlushClient, PublishClient, RequestClient, session};
+use crate::nats::{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, JsRequestMessage};
 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> {
+) -> Result<ForkSessionResponse>
+where
+    <J::Consumer as trogon_nats::jetstream::JetStreamConsumer>::Message: JsRequestMessage,
+{
     let start = bridge.clock.now();
 
     info!(session_id = %args.session_id, "Fork session request");
@@ -28,17 +35,16 @@ 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 = bridge
+        .session_request::<ForkSessionRequest, ForkSessionResponse>(
+            &subject,
+            &args,
+            session_id.as_str(),
+        )
+        .await;
 
     if let Ok(ref response) = result {
         Span::current().record("new_session_id", response.session_id.to_string().as_str());