feat(trogon-nats, acp-nats): add JetStream integration foundation (#65)

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

Author: yordis

rsworkspace/Cargo.lock

@@ -19,7 +19,7 @@ agent-client-protocol = { workspace = true, features = [
     "unstable_session_resume",
     "unstable_session_usage",
 ] }
-async-nats = { workspace = true }
+async-nats = { workspace = true, features = ["jetstream"] }
 async-trait = { workspace = true }
 bytes = { workspace = true }
 futures = { workspace = true }

rsworkspace/crates/acp-nats-agent/Cargo.toml

@@ -16,6 +16,8 @@ use futures::future::LocalBoxFuture;
 use std::rc::Rc;
 use std::time::Duration;
 use tracing::{info, warn};
+#[cfg(not(coverage))]
+use trogon_nats::jetstream::JsMessage;
 use trogon_nats::{FlushClient, PublishClient, RequestClient, SubscribeClient};
 
 pub enum ConnectionError {
@@ -313,6 +315,137 @@ where
         .map_err(DispatchError::NotificationHandler)
 }
 
+/// JetStream-aware dispatch: receives a `JsMessage`, dispatches to the agent,
+/// and signals ack/term based on the outcome.
+///
+/// - Unknown subject → `term()` (no redelivery)
+/// - Deserialize failure → reply error + `term()` (bad payload won't fix itself)
+/// - Handler success → reply + `ack()`
+/// - Handler error → reply error + `ack()` (application-level error, not transient)
+/// - Notification handler → `ack()`
+#[cfg(not(coverage))]
+#[allow(dead_code)] // Will be used when JetStream serve path is wired up
+async fn dispatch_js_message<N: PublishClient + FlushClient, A: Agent>(
+    js_msg: &JsMessage,
+    agent: &A,
+    nats: &N,
+) {
+    let msg = js_msg.message();
+    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(msg, nats, |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(_) => {
+            if let Err(e) = js_msg.ack().await {
+                warn!(subject, error = %e, "Failed to ack after 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-agent/src/connection.rs

@@ -19,7 +19,7 @@ agent-client-protocol = { workspace = true, features = [
     "unstable_session_usage",
 ] }
 opentelemetry = { workspace = true }
-async-nats = { workspace = true }
+async-nats = { workspace = true, features = ["jetstream"] }
 async-trait = { workspace = true }
 bytes = { workspace = true }
 futures = { workspace = true }

rsworkspace/crates/acp-nats/Cargo.toml

@@ -10,8 +10,8 @@ use trogon_std::time::GetElapsed;
     skip(bridge, args),
     fields(method_id = %args.method_id)
 )]
-pub async fn handle<N: RequestClient, C: GetElapsed>(
-    bridge: &Bridge<N, C>,
+pub async fn handle<N: RequestClient, C: GetElapsed, J>(
+    bridge: &Bridge<N, C, J>,
     args: AuthenticateRequest,
 ) -> Result<AuthenticateResponse> {
     let start = bridge.clock.now();

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

@@ -21,6 +21,9 @@ 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;
 
 use super::{
@@ -31,8 +34,10 @@ use super::{
 
 use crate::constants::SESSION_READY_DELAY;
 
-pub struct Bridge<N, C: GetElapsed> {
+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,
     pub(crate) metrics: Metrics,
@@ -51,6 +56,30 @@ impl<N, C: GetElapsed> Bridge<N, C> {
     ) -> Self {
         Self {
             nats,
+            js: None,
+            clock,
+            config,
+            metrics: Metrics::new(meter),
+            notification_sender,
+            pending_session_prompt_responses: PendingSessionPromptResponseWaiters::new(),
+            background_tasks: RefCell::new(Vec::new()),
+        }
+    }
+}
+
+impl<N, C: GetElapsed, J> Bridge<N, C, J> {
+    #[cfg(not(coverage))]
+    pub fn with_jetstream(
+        nats: N,
+        js: J,
+        clock: C,
+        meter: &Meter,
+        config: Config,
+        notification_sender: mpsc::Sender<SessionNotification>,
+    ) -> Self {
+        Self {
+            nats,
+            js: Some(js),
             clock,
             config,
             metrics: Metrics::new(meter),
@@ -64,6 +93,12 @@ impl<N, C: GetElapsed> Bridge<N, C> {
         &self.nats
     }
 
+    #[cfg(not(coverage))]
+    #[allow(dead_code)]
+    pub(crate) fn js(&self) -> Option<&J> {
+        self.js.as_ref()
+    }
+
     pub(crate) fn spawn_background(&self, task: JoinHandle<()>) {
         self.background_tasks.borrow_mut().push(task);
     }
@@ -76,7 +111,7 @@ impl<N, C: GetElapsed> Bridge<N, C> {
     }
 }
 
-impl<N: PublishClient + FlushClient + Clone + Send + 'static, C: GetElapsed> Bridge<N, C> {
+impl<N: PublishClient + FlushClient + Clone + Send + 'static, C: GetElapsed, J> Bridge<N, C, J> {
     pub(crate) fn schedule_session_ready(&self, session_id: SessionId) {
         let nats = self.nats.clone();
         let prefix = self.config.acp_prefix().to_string();
@@ -119,8 +154,8 @@ async fn publish_session_ready<N: PublishClient + FlushClient>(
 }
 
 #[async_trait::async_trait(?Send)]
-impl<N: RequestClient + PublishClient + SubscribeClient + FlushClient, C: GetElapsed> Agent
-    for Bridge<N, C>
+impl<N: RequestClient + PublishClient + SubscribeClient + FlushClient, C: GetElapsed, J> Agent
+    for Bridge<N, C, J>
 {
     async fn initialize(&self, args: InitializeRequest) -> Result<InitializeResponse> {
         initialize::handle(self, args).await