fix(acp-nats): replace unimplemented!() in session_update test mock

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

Author: yordis

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

@@ -7,6 +7,7 @@ use crate::nats::{
     headers_with_trace_context, parse_client_subject,
 };
 use agent_client_protocol::{Client, Error, JsonRpcMessage, RequestId, Response};
+use async_nats::Message;
 use bytes::Bytes;
 use futures::StreamExt;
 use std::cell::Cell;
@@ -66,90 +67,105 @@ pub async fn run<
     let max_concurrent = bridge.config.max_concurrent_client_tasks();
 
     while let Some(msg) = subscriber.next().await {
-        let subject = msg.subject.to_string();
-
-        // Validate subject before backpressure so unrecognised methods always
-        // get InvalidParams, not a misleading "Bridge overloaded" error.
-        let parsed = match parse_client_subject(&subject) {
-            Some(parsed) => parsed,
-            None => {
-                let request_id = extract_request_id(&msg.payload);
-                warn!(subject = %subject, "Failed to parse client subject");
-                if let Some(reply_to) = msg.reply.as_ref().map(|reply| reply.to_string()) {
-                    let bytes = error_response_bytes(
-                        request_id,
-                        Error::invalid_params().data("Invalid client subject"),
-                    );
-                    let headers = headers_with_trace_context();
-                    if let Err(e) = nats.publish_with_headers(reply_to, headers, bytes).await {
-                        error!(error = %e, "Failed to publish invalid subject response");
-                    }
-                    if let Err(error) = nats.flush().await {
-                        warn!(error = %error, "Failed to flush invalid subject response");
-                    }
-                }
-                continue;
-            }
-        };
+        process_message(msg, &nats, client.clone(), bridge.clone(), &in_flight, max_concurrent)
+            .await;
+    }
 
-        if in_flight.get() >= max_concurrent {
-            let request_id = extract_request_id(&msg.payload);
-            warn!(
-                in_flight = in_flight.get(),
-                method = ?parsed.method,
-                subject = %subject,
-                "Client task backpressure — rejecting message"
-            );
-            bridge
-                .metrics
-                .record_error("client", "client_backpressure_rejected");
+    info!("Client proxy subscriber ended");
+}
 
+async fn process_message<
+    N: SubscribeClient + RequestClient + PublishClient + FlushClient,
+    Cl: Client + 'static,
+    C: GetElapsed + 'static,
+>(
+    msg: Message,
+    nats: &N,
+    client: Rc<Cl>,
+    bridge: Rc<Bridge<N, C>>,
+    in_flight: &Rc<Cell<usize>>,
+    max_concurrent: usize,
+) {
+    let subject = msg.subject.to_string();
+
+    // Validate subject before backpressure so unrecognised methods always
+    // get InvalidParams, not a misleading "Bridge overloaded" error.
+    let parsed = match parse_client_subject(&subject) {
+        Some(parsed) => parsed,
+        None => {
+            let request_id = extract_request_id(&msg.payload);
+            warn!(subject = %subject, "Failed to parse client subject");
             if let Some(reply_to) = msg.reply.as_ref().map(|reply| reply.to_string()) {
                 let bytes = error_response_bytes(
                     request_id,
-                    Error::internal_error()
-                        .data("Bridge overloaded - too many concurrent requests"),
+                    Error::invalid_params().data("Invalid client subject"),
                 );
                 let headers = headers_with_trace_context();
                 if let Err(e) = nats.publish_with_headers(reply_to, headers, bytes).await {
-                    error!(error = %e, "Failed to publish backpressure response");
+                    error!(error = %e, "Failed to publish invalid subject response");
                 }
                 if let Err(error) = nats.flush().await {
-                    warn!(error = %error, "Failed to flush backpressure response");
+                    warn!(error = %error, "Failed to flush invalid subject response");
                 }
-            } else {
-                warn!(
-                    subject = %subject,
-                    in_flight = in_flight.get(),
-                    "No reply_to on request; dropping due to backpressure"
-                );
             }
-            continue;
+            return;
         }
+    };
 
-        let reply = msg.reply.as_ref().map(|reply| reply.to_string());
-        let payload = msg.payload.clone();
-        let nats = nats.clone();
-        let client = client.clone();
-
-        let bridge_clone = bridge.clone();
-        let in_flight_guard = InFlightSlotGuard::new(in_flight.clone());
-        tokio::task::spawn_local(async move {
-            let _in_flight_guard = in_flight_guard;
-            handle_client_request(
-                &subject,
-                parsed,
-                payload,
-                reply,
-                &nats,
-                client.as_ref(),
-                bridge_clone.as_ref(),
-            )
-            .await;
-        });
+    if in_flight.get() >= max_concurrent {
+        let request_id = extract_request_id(&msg.payload);
+        warn!(
+            in_flight = in_flight.get(),
+            method = ?parsed.method,
+            subject = %subject,
+            "Client task backpressure — rejecting message"
+        );
+        bridge
+            .metrics
+            .record_error("client", "client_backpressure_rejected");
+
+        if let Some(reply_to) = msg.reply.as_ref().map(|reply| reply.to_string()) {
+            let bytes = error_response_bytes(
+                request_id,
+                Error::internal_error()
+                    .data("Bridge overloaded - too many concurrent requests"),
+            );
+            let headers = headers_with_trace_context();
+            if let Err(e) = nats.publish_with_headers(reply_to, headers, bytes).await {
+                error!(error = %e, "Failed to publish backpressure response");
+            }
+            if let Err(error) = nats.flush().await {
+                warn!(error = %error, "Failed to flush backpressure response");
+            }
+        } else {
+            warn!(
+                subject = %subject,
+                in_flight = in_flight.get(),
+                "No reply_to on request; dropping due to backpressure"
+            );
+        }
+        return;
     }
 
-    info!("Client proxy subscriber ended");
+    let reply = msg.reply.as_ref().map(|reply| reply.to_string());
+    let payload = msg.payload.clone();
+    let nats = nats.clone();
+
+    let bridge_clone = bridge.clone();
+    let in_flight_guard = InFlightSlotGuard::new(in_flight.clone());
+    tokio::task::spawn_local(async move {
+        let _in_flight_guard = in_flight_guard;
+        handle_client_request(
+            &subject,
+            parsed,
+            payload,
+            reply,
+            &nats,
+            client.as_ref(),
+            bridge_clone.as_ref(),
+        )
+        .await;
+    });
 }
 
 #[instrument(skip(payload, reply, _nats, client, _bridge), fields(subject = %subject, session_id = tracing::field::Empty))]
@@ -221,15 +237,33 @@ mod tests {
         }
     }
 
-    #[tokio::test]
-    async fn run_returns_early_when_subscribe_fails() {
-        let nats = MockNatsClient::new();
-        let bridge = Rc::new(Bridge::new(
-            nats.clone(),
+    fn make_msg(subject: &str, payload: &[u8], reply: Option<&str>) -> async_nats::Message {
+        async_nats::Message {
+            subject: subject.into(),
+            reply: reply.map(|r| r.into()),
+            payload: payload.to_vec().into(),
+            headers: None,
+            length: payload.len(),
+            status: None,
+            description: None,
+        }
+    }
+
+    fn make_bridge(
+        nats: MockNatsClient,
+    ) -> Rc<Bridge<MockNatsClient, SystemClock>> {
+        Rc::new(Bridge::new(
+            nats,
             SystemClock,
             &opentelemetry::global::meter("acp-nats-test"),
             crate::config::Config::for_test("acp"),
-        ));
+        ))
+    }
+
+    #[tokio::test]
+    async fn run_returns_early_when_subscribe_fails() {
+        let nats = MockNatsClient::new();
+        let bridge = make_bridge(nats.clone());
         let client = Rc::new(MockClient::new());
 
         run(nats, client, bridge).await;
@@ -304,4 +338,89 @@ mod tests {
 
         assert_eq!(client.notifications.borrow().len(), 1);
     }
+
+    #[tokio::test]
+    async fn process_message_invalid_subject_no_reply_does_not_publish() {
+        let nats = MockNatsClient::new();
+        let bridge = make_bridge(nats.clone());
+        let client = Rc::new(MockClient::new());
+        let in_flight = Rc::new(Cell::new(0usize));
+
+        let msg = make_msg("acp.sess.unknown.method", b"{}", None);
+        process_message(msg, &nats, client, bridge, &in_flight, 256).await;
+
+        assert!(nats.published_messages().is_empty());
+    }
+
+    #[tokio::test]
+    async fn process_message_invalid_subject_with_reply_publishes_error() {
+        let nats = MockNatsClient::new();
+        let bridge = make_bridge(nats.clone());
+        let client = Rc::new(MockClient::new());
+        let in_flight = Rc::new(Cell::new(0usize));
+
+        let msg = make_msg("acp.sess.unknown.method", b"{}", Some("_INBOX.reply"));
+        process_message(msg, &nats, client, bridge, &in_flight, 256).await;
+
+        let published = nats.published_messages();
+        assert_eq!(published, vec!["_INBOX.reply"]);
+    }
+
+    #[tokio::test]
+    async fn process_message_backpressure_no_reply_does_not_publish() {
+        let nats = MockNatsClient::new();
+        let bridge = make_bridge(nats.clone());
+        let client = Rc::new(MockClient::new());
+        let in_flight = Rc::new(Cell::new(1usize));
+
+        let msg = make_msg("acp.sess1.client.session.update", b"{}", None);
+        process_message(msg, &nats, client, bridge, &in_flight, 1).await;
+
+        assert!(nats.published_messages().is_empty());
+    }
+
+    #[tokio::test]
+    async fn process_message_backpressure_with_reply_publishes_error() {
+        let nats = MockNatsClient::new();
+        let bridge = make_bridge(nats.clone());
+        let client = Rc::new(MockClient::new());
+        let in_flight = Rc::new(Cell::new(1usize));
+
+        let msg = make_msg(
+            "acp.sess1.client.session.update",
+            b"{}",
+            Some("_INBOX.reply"),
+        );
+        process_message(msg, &nats, client, bridge, &in_flight, 1).await;
+
+        let published = nats.published_messages();
+        assert_eq!(published, vec!["_INBOX.reply"]);
+    }
+
+    #[tokio::test]
+    async fn process_message_valid_dispatch_spawns_task() {
+        let local = tokio::task::LocalSet::new();
+        local
+            .run_until(async {
+                let nats = MockNatsClient::new();
+                let bridge = make_bridge(nats.clone());
+                let client = Rc::new(MockClient::new());
+                let in_flight = Rc::new(Cell::new(0usize));
+
+                let notification = SessionNotification::new(
+                    "sess1",
+                    SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::from("hi"))),
+                );
+                let payload = serde_json::to_vec(&notification).unwrap();
+
+                let msg = make_msg("acp.sess1.client.session.update", &payload, None);
+                process_message(msg, &nats, client.clone(), bridge, &in_flight, 256).await;
+
+                // Yield to allow the spawned local task to run.
+                tokio::task::yield_now().await;
+
+                assert_eq!(client.notifications.borrow().len(), 1);
+            })
+            .await;
+    }
 }

rsworkspace/crates/acp-nats/src/client/session_update.rs

@@ -95,7 +95,10 @@ mod tests {
             &self,
             _: RequestPermissionRequest,
         ) -> Result<RequestPermissionResponse, agent_client_protocol::Error> {
-            unimplemented!()
+            Err(agent_client_protocol::Error::new(
+                -32603,
+                "not implemented in test mock",
+            ))
         }
     }