refactor(acp-nats, acp-nats-agent): narrow JetStream scope and add ExtStreamPolicy (#66)

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

Author: yordis

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

@@ -16,8 +16,6 @@ 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 {
@@ -315,137 +313,6 @@ 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/src/constants.rs

@@ -32,4 +32,3 @@ pub const CONTENT_TYPE_PLAIN: &str = "text/plain";
 
 pub const SESSION_ID_HEADER: &str = "X-Session-Id";
 pub const CAUSATION_ID_HEADER: &str = "X-Causation-Id";
-pub const ENV_JETSTREAM_ENABLED: &str = "ACP_JETSTREAM_ENABLED";

rsworkspace/crates/acp-nats/src/jetstream/consumers.rs

@@ -21,11 +21,11 @@ pub fn prompt_response_consumer(prefix: &str, session_id: &str, req_id: &str) ->
     }
 }
 
-/// Consumer for the COMMANDS stream that receives all commands.
+/// Observer consumer for the COMMANDS stream.
 ///
-/// No filter needed — the COMMANDS stream already contains only command subjects.
-/// The stream-level subject list acts as the filter.
-pub fn runner_commands_consumer() -> Config {
+/// Acks messages for audit persistence. No filter needed — the stream-level
+/// subject list already scopes to session-scoped commands only.
+pub fn commands_observer() -> Config {
     Config {
         deliver_policy: DeliverPolicy::All,
         ack_policy: AckPolicy::Explicit,
@@ -65,15 +65,15 @@ mod tests {
     }
 
     #[test]
-    fn runner_commands_consumer_delivers_all() {
-        let config = runner_commands_consumer();
+    fn commands_observer_delivers_all() {
+        let config = commands_observer();
         assert_eq!(config.deliver_policy, DeliverPolicy::All);
         assert_eq!(config.ack_policy, AckPolicy::Explicit);
     }
 
     #[test]
-    fn runner_commands_consumer_no_filter() {
-        let config = runner_commands_consumer();
+    fn commands_observer_no_filter() {
+        let config = commands_observer();
         assert_eq!(config.filter_subject, String::new());
     }
 

rsworkspace/crates/acp-nats/src/jetstream/ext_policy.rs

@@ -0,0 +1,64 @@
+use crate::ext_method_name::ExtMethodName;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum ExtPersistence {
+    Stream,
+    Ephemeral,
+}
+
+pub trait ExtStreamPolicy: Send + Sync {
+    fn persistence(&self, method: &ExtMethodName) -> ExtPersistence;
+}
+
+pub struct DefaultExtStreamPolicy;
+
+impl ExtStreamPolicy for DefaultExtStreamPolicy {
+    fn persistence(&self, _method: &ExtMethodName) -> ExtPersistence {
+        ExtPersistence::Stream
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn default_policy_persists_everything() {
+        let policy = DefaultExtStreamPolicy;
+        let method = ExtMethodName::new("tool_call").unwrap();
+        assert_eq!(policy.persistence(&method), ExtPersistence::Stream);
+    }
+
+    #[test]
+    fn custom_policy() {
+        struct SelectivePolicy;
+        impl ExtStreamPolicy for SelectivePolicy {
+            fn persistence(&self, method: &ExtMethodName) -> ExtPersistence {
+                match method.as_str() {
+                    "heartbeat" | "ping" => ExtPersistence::Ephemeral,
+                    _ => ExtPersistence::Stream,
+                }
+            }
+        }
+
+        let policy = SelectivePolicy;
+        assert_eq!(
+            policy.persistence(&ExtMethodName::new("heartbeat").unwrap()),
+            ExtPersistence::Ephemeral
+        );
+        assert_eq!(
+            policy.persistence(&ExtMethodName::new("ping").unwrap()),
+            ExtPersistence::Ephemeral
+        );
+        assert_eq!(
+            policy.persistence(&ExtMethodName::new("tool_call").unwrap()),
+            ExtPersistence::Stream
+        );
+    }
+
+    #[test]
+    fn ext_persistence_debug() {
+        assert_eq!(format!("{:?}", ExtPersistence::Stream), "Stream");
+        assert_eq!(format!("{:?}", ExtPersistence::Ephemeral), "Ephemeral");
+    }
+}

rsworkspace/crates/acp-nats/src/jetstream/mod.rs

@@ -1,3 +1,4 @@
 pub mod consumers;
+pub mod ext_policy;
 pub mod provision;
 pub mod streams;

rsworkspace/crates/acp-nats/src/jetstream/streams.rs

@@ -13,7 +13,6 @@ pub fn commands_config(prefix: &str) -> Config {
     Config {
         name: stream_name(prefix, "COMMANDS"),
         subjects: vec![
-            format!("{prefix}.agent.>"),
             format!("{prefix}.session.*.agent.prompt"),
             format!("{prefix}.session.*.agent.cancel"),
             format!("{prefix}.session.*.agent.load"),
@@ -23,7 +22,6 @@ pub fn commands_config(prefix: &str) -> Config {
             format!("{prefix}.session.*.agent.fork"),
             format!("{prefix}.session.*.agent.resume"),
             format!("{prefix}.session.*.agent.close"),
-            format!("{prefix}.session.*.agent.ext.>"),
         ],
         storage: StorageType::File,
         retention: RetentionPolicy::Limits,
@@ -101,9 +99,9 @@ mod tests {
     }
 
     #[test]
-    fn commands_subjects_include_global_and_session() {
+    fn commands_subjects_are_session_scoped_only() {
         let config = commands_config("acp");
-        assert!(config.subjects.contains(&"acp.agent.>".to_string()));
+        assert!(!config.subjects.contains(&"acp.agent.>".to_string()));
         assert!(
             config
                 .subjects
@@ -117,10 +115,16 @@ mod tests {
         assert!(
             config
                 .subjects
-                .contains(&"acp.session.*.agent.ext.>".to_string())
+                .contains(&"acp.session.*.agent.close".to_string())
         );
     }
 
+    #[test]
+    fn commands_excludes_ext_subjects() {
+        let config = commands_config("acp");
+        assert!(!config.subjects.iter().any(|s| s.contains("ext")));
+    }
+
     #[test]
     fn responses_subjects() {
         let config = responses_config("acp");