refactor(acp-nats): typed subjects enforce correct transport at compile time

Introduce six subject newtypes: RequestSubject, PublishSubject,
JetStreamSubject, AgentResponseSubject, SubscribeSubject, and
ClientRequestSubject. Every subject builder returns the type matching
its transport. Wrapper functions only accept the correct type.

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

Author: yordis

rsworkspace/Cargo.lock

@@ -12,6 +12,7 @@ use agent_client_protocol::{
 };
 use async_nats::Message;
 use async_nats::jetstream::AckKind;
+use async_nats::subject::ToSubject;
 #[cfg(test)]
 use bytes::Bytes;
 use futures::StreamExt;
@@ -112,8 +113,8 @@ where
         impl std::future::Future<Output = Result<(), ConnectionError>>,
     )
     where
-        J: JetStreamConsumerFactory + 'static,
-        <J::Consumer as trogon_nats::jetstream::JetStreamConsumer>::Message: JsDispatchMessage,
+        J: JetStreamGetStream + 'static,
+        <<J::Stream as trogon_nats::jetstream::JetStreamCreateConsumer>::Consumer as trogon_nats::jetstream::JetStreamConsumer>::Message: JsDispatchMessage,
     {
         let nats_for_serve = nats.clone();
         let nats_for_js = nats.clone();
@@ -422,7 +423,8 @@ where
 }
 
 use trogon_nats::jetstream::{
-    JetStreamConsumer as _, JetStreamConsumerFactory, JsAckWith, JsDispatchMessage,
+    JetStreamConsumer as _, JetStreamCreateConsumer as _, JetStreamGetStream, JsAckWith,
+    JsDispatchMessage,
 };
 
 async fn handle_request_with_keepalive<N, Resp, ReqT, F, M>(
@@ -484,17 +486,22 @@ async fn serve_js<N, J, A>(
 ) -> Result<(), ConnectionError>
 where
     N: PublishClient + FlushClient + Clone + 'static,
-    J: JetStreamConsumerFactory + 'static,
-    <J::Consumer as trogon_nats::jetstream::JetStreamConsumer>::Message: JsDispatchMessage,
+    J: JetStreamGetStream + 'static,
+    <<J::Stream as trogon_nats::jetstream::JetStreamCreateConsumer>::Consumer as trogon_nats::jetstream::JetStreamConsumer>::Message: JsDispatchMessage,
     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)
+    let stream = js
+        .get_stream(&stream_name)
+        .await
+        .map_err(|e| ConnectionError::JetStream(Box::new(e)))?;
+
+    let consumer = stream
+        .create_consumer(config)
         .await
         .map_err(|e| ConnectionError::JetStream(Box::new(e)))?;
 
@@ -581,7 +588,7 @@ async fn dispatch_js_message<N: PublishClient + FlushClient, A: Agent, M: JsDisp
     let inner = js_msg.message();
     let msg = Message {
         subject: subject.as_str().into(),
-        reply: reply_subject.as_deref().map(|s| s.into()),
+        reply: reply_subject.map(|s| s.to_subject()),
         payload: inner.payload.clone(),
         headers: inner.headers.clone(),
         status: None,

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

@@ -21,7 +21,7 @@ use opentelemetry::metrics::Meter;
 use tokio::sync::mpsc;
 use tokio::task::JoinHandle;
 use tracing::{info, warn};
-use trogon_nats::jetstream::{JetStreamConsumerFactory, JetStreamPublisher, JsRequestMessage};
+use trogon_nats::jetstream::{JetStreamGetStream, JetStreamPublisher, JsRequestMessage};
 use trogon_std::time::GetElapsed;
 
 use super::{
@@ -150,14 +150,14 @@ async fn publish_session_ready<N: PublishClient + FlushClient>(
 impl<
     N: RequestClient + PublishClient + FlushClient,
     C: GetElapsed,
-    J: JetStreamPublisher + JetStreamConsumerFactory,
+    J: JetStreamPublisher + JetStreamGetStream,
 > Bridge<N, C, J>
 where
-    <J::Consumer as trogon_nats::jetstream::JetStreamConsumer>::Message: JsRequestMessage,
+    <<J::Stream as trogon_nats::jetstream::JetStreamCreateConsumer>::Consumer as trogon_nats::jetstream::JetStreamConsumer>::Message: JsRequestMessage,
 {
     pub(crate) async fn session_request<Req, Res>(
         &self,
-        subject: &str,
+        subject: &crate::nats::JetStreamSubject,
         args: &Req,
         session_id: &str,
     ) -> Result<Res>
@@ -172,7 +172,7 @@ where
                 let req_id = uuid::Uuid::new_v4().to_string();
                 js_request::js_request::<J, _, Res, _>(
                     js,
-                    subject,
+                    &subject.0,
                     args,
                     &trogon_std::StdJsonSerialize,
                     self.config.acp_prefix(),
@@ -182,9 +182,9 @@ where
                 )
                 .await
             }
-            None => nats::request_with_timeout::<N, Req, Res>(
+            None => trogon_nats::request_with_timeout::<N, Req, Res>(
                 self.nats(),
-                subject,
+                &subject.0,
                 args,
                 self.config.operation_timeout,
             )
@@ -198,10 +198,10 @@ where
 impl<
     N: RequestClient + PublishClient + SubscribeClient + FlushClient,
     C: GetElapsed,
-    J: JetStreamPublisher + JetStreamConsumerFactory,
+    J: JetStreamPublisher + JetStreamGetStream,
 > Agent for Bridge<N, C, J>
 where
-    <J::Consumer as trogon_nats::jetstream::JetStreamConsumer>::Message: JsRequestMessage,
+    <<J::Stream as trogon_nats::jetstream::JetStreamCreateConsumer>::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/bridge.rs

@@ -3,7 +3,7 @@ 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_nats::jetstream::{JetStreamGetStream, JetStreamPublisher, JsRequestMessage};
 use trogon_std::time::GetElapsed;
 
 #[instrument(
@@ -14,13 +14,13 @@ use trogon_std::time::GetElapsed;
 pub async fn handle<
     N: RequestClient + PublishClient + FlushClient,
     C: GetElapsed,
-    J: JetStreamPublisher + JetStreamConsumerFactory,
+    J: JetStreamPublisher + JetStreamGetStream,
 >(
     bridge: &Bridge<N, C, J>,
     args: CloseSessionRequest,
 ) -> Result<CloseSessionResponse>
 where
-    <J::Consumer as trogon_nats::jetstream::JetStreamConsumer>::Message: JsRequestMessage,
+    <<J::Stream as trogon_nats::jetstream::JetStreamCreateConsumer>::Consumer as trogon_nats::jetstream::JetStreamConsumer>::Message: JsRequestMessage,
 {
     let start = bridge.clock.now();
 

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

@@ -38,7 +38,7 @@ pub async fn handle<N: PublishClient + FlushClient, C: GetElapsed, J>(
         )
     })?;
 
-    let subject = agent::ext(bridge.config.acp_prefix(), method_name.as_str());
+    let subject = agent::ext_notify(bridge.config.acp_prefix(), method_name.as_str());
 
     let publish_result = nats::publish(
         bridge.nats(),

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

@@ -3,7 +3,7 @@ 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_nats::jetstream::{JetStreamGetStream, JetStreamPublisher, JsRequestMessage};
 use trogon_std::time::GetElapsed;
 
 #[instrument(
@@ -14,13 +14,13 @@ use trogon_std::time::GetElapsed;
 pub async fn handle<
     N: RequestClient + PublishClient + FlushClient,
     C: GetElapsed,
-    J: JetStreamPublisher + JetStreamConsumerFactory,
+    J: JetStreamPublisher + JetStreamGetStream,
 >(
     bridge: &Bridge<N, C, J>,
     args: ForkSessionRequest,
 ) -> Result<ForkSessionResponse>
 where
-    <J::Consumer as trogon_nats::jetstream::JetStreamConsumer>::Message: JsRequestMessage,
+    <<J::Stream as trogon_nats::jetstream::JetStreamCreateConsumer>::Consumer as trogon_nats::jetstream::JetStreamConsumer>::Message: JsRequestMessage,
 {
     let start = bridge.clock.now();
 

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

@@ -7,8 +7,8 @@ use std::time::Duration;
 use tokio::time::timeout;
 use trogon_nats::REQ_ID_HEADER;
 use trogon_nats::jetstream::{
-    JetStreamConsumer as _, JetStreamConsumerFactory, JetStreamPublisher, JsAck as _,
-    JsAckWith as _, JsMessageRef as _, JsRequestMessage,
+    JetStreamConsumer as _, JetStreamCreateConsumer as _, JetStreamGetStream, JetStreamPublisher,
+    JsAck as _, JsAckWith as _, JsMessageRef as _, JsRequestMessage,
 };
 use trogon_std::JsonSerialize;
 
@@ -27,8 +27,8 @@ pub async fn js_request<J, Req, Res, S>(
     operation_timeout: Duration,
 ) -> agent_client_protocol::Result<Res>
 where
-    J: JetStreamPublisher + JetStreamConsumerFactory,
-    <J::Consumer as trogon_nats::jetstream::JetStreamConsumer>::Message: JsRequestMessage,
+    J: JetStreamPublisher + JetStreamGetStream,
+    <<J::Stream as trogon_nats::jetstream::JetStreamCreateConsumer>::Consumer as trogon_nats::jetstream::JetStreamConsumer>::Message: JsRequestMessage,
     Req: serde::Serialize,
     Res: DeserializeOwned,
     S: JsonSerialize,
@@ -37,22 +37,22 @@ where
     // runner responds before we start consuming. DeliverAll replays from stream start.
     let responses_stream = streams::responses_stream_name(prefix);
     let resp_config = consumers::response_consumer(prefix, session_id, req_id);
-    let resp_consumer: J::Consumer = js
-        .create_consumer(&responses_stream, resp_config)
+    let stream = js
+        .get_stream(&responses_stream)
         .await
-        .map_err(|e| {
-            Error::new(
-                ErrorCode::InternalError.into(),
-                format!("create response consumer: {e}"),
-            )
-        })?;
-    let mut resp_messages: <J::Consumer as trogon_nats::jetstream::JetStreamConsumer>::Messages =
-        resp_consumer.messages().await.map_err(|e| {
-            Error::new(
-                ErrorCode::InternalError.into(),
-                format!("response messages: {e}"),
-            )
-        })?;
+        .map_err(|e| Error::new(ErrorCode::InternalError.into(), format!("get stream: {e}")))?;
+    let resp_consumer = stream.create_consumer(resp_config).await.map_err(|e| {
+        Error::new(
+            ErrorCode::InternalError.into(),
+            format!("create response consumer: {e}"),
+        )
+    })?;
+    let mut resp_messages = resp_consumer.messages().await.map_err(|e| {
+        Error::new(
+            ErrorCode::InternalError.into(),
+            format!("response messages: {e}"),
+        )
+    })?;
 
     let payload_bytes = serializer
         .to_vec(request)
@@ -62,9 +62,11 @@ where
     headers.insert(REQ_ID_HEADER, req_id);
     headers.insert(SESSION_ID_HEADER, session_id);
 
-    js.js_publish_with_headers(subject.to_string(), headers, Bytes::from(payload_bytes))
+    js.publish_with_headers(subject.to_string(), headers, Bytes::from(payload_bytes))
+        .await
+        .map_err(|e| Error::new(ErrorCode::InternalError.into(), format!("js publish: {e}")))?
         .await
-        .map_err(|e| Error::new(ErrorCode::InternalError.into(), format!("js publish: {e}")))?;
+        .map_err(|e| Error::new(ErrorCode::InternalError.into(), format!("js ack: {e}")))?;
 
     match timeout(operation_timeout, resp_messages.next()).await {
         Ok(Some(Ok(js_msg))) => {

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

@@ -3,7 +3,7 @@ use crate::nats::{FlushClient, PublishClient, RequestClient, session};
 use crate::session_id::AcpSessionId;
 use agent_client_protocol::{Error, ErrorCode, LoadSessionRequest, LoadSessionResponse, Result};
 use tracing::{info, instrument};
-use trogon_nats::jetstream::{JetStreamConsumerFactory, JetStreamPublisher, JsRequestMessage};
+use trogon_nats::jetstream::{JetStreamGetStream, JetStreamPublisher, JsRequestMessage};
 use trogon_std::time::GetElapsed;
 
 #[instrument(
@@ -14,13 +14,13 @@ use trogon_std::time::GetElapsed;
 pub async fn handle<
     N: RequestClient + PublishClient + FlushClient,
     C: GetElapsed,
-    J: JetStreamPublisher + JetStreamConsumerFactory,
+    J: JetStreamPublisher + JetStreamGetStream,
 >(
     bridge: &Bridge<N, C, J>,
     args: LoadSessionRequest,
 ) -> Result<LoadSessionResponse>
 where
-    <J::Consumer as trogon_nats::jetstream::JetStreamConsumer>::Message: JsRequestMessage,
+    <<J::Stream as trogon_nats::jetstream::JetStreamCreateConsumer>::Consumer as trogon_nats::jetstream::JetStreamConsumer>::Message: JsRequestMessage,
 {
     let start = bridge.clock.now();