feat(trogon-nats): propagate X-Req-Id header across NATS requests

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

Author: yordis

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

@@ -256,17 +256,29 @@ async fn reply<N: PublishClient + FlushClient, T: serde::Serialize>(
     nats: &N,
     reply_to: &str,
     value: &T,
+    req_id: Option<&str>,
 ) -> Result<(), DispatchError> {
-    trogon_nats::publish(
-        nats,
-        reply_to,
-        value,
-        trogon_nats::PublishOptions::builder()
-            .flush_policy(trogon_nats::FlushPolicy::standard())
-            .build(),
-    )
-    .await
-    .map_err(DispatchError::Reply)
+    let mut headers = trogon_nats::headers_with_trace_context();
+    if let Some(id) = req_id {
+        headers.insert(trogon_nats::REQ_ID_HEADER, id);
+    }
+
+    let payload = serde_json::to_vec(value)
+        .map_err(|e| DispatchError::Reply(trogon_nats::NatsError::Serialize(e)))?;
+
+    nats.publish_with_headers(reply_to.to_string(), headers, payload.into())
+        .await
+        .map_err(|e| {
+            DispatchError::Reply(trogon_nats::NatsError::PublishOperation(
+                trogon_nats::PublishOperationError(e.to_string()),
+            ))
+        })?;
+
+    nats.flush().await.map_err(|e| {
+        DispatchError::Reply(trogon_nats::NatsError::PublishOperation(
+            trogon_nats::PublishOperationError(e.to_string()),
+        ))
+    })
 }
 
 async fn handle_request<N, Resp, ReqT, F>(
@@ -281,6 +293,7 @@ where
     Resp: serde::Serialize,
 {
     let reply_to = msg.reply.as_deref().ok_or(DispatchError::NoReplySubject)?;
+    let req_id = trogon_nats::extract_req_id(msg.headers.as_ref());
 
     let request: ReqT = match serde_json::from_slice(&msg.payload) {
         Ok(req) => req,
@@ -289,14 +302,14 @@ where
                 agent_client_protocol::ErrorCode::InvalidParams.into(),
                 format!("Failed to deserialize request: {}", e),
             );
-            let _ = reply(nats, reply_to, &error).await;
+            let _ = reply(nats, reply_to, &error, req_id).await;
             return Err(DispatchError::DeserializeRequest(e));
         }
     };
 
     match handler(request).await {
-        Ok(resp) => reply(nats, reply_to, &resp).await,
-        Err(err) => reply(nats, reply_to, &err).await,
+        Ok(resp) => reply(nats, reply_to, &resp, req_id).await,
+        Err(err) => reply(nats, reply_to, &err, req_id).await,
     }
 }
 

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

@@ -11,7 +11,7 @@ use crate::agent::Bridge;
 use crate::nats::{FlushClient, PublishClient, RequestClient, SubscribeClient, agent};
 use crate::session_id::AcpSessionId;
 
-pub const REQ_ID_HEADER: &str = "X-Req-Id";
+pub use trogon_nats::REQ_ID_HEADER;
 
 #[instrument(
     name = "acp.session.prompt",

rsworkspace/crates/trogon-nats/Cargo.toml

@@ -17,6 +17,7 @@ tokio = { workspace = true, features = ["time"] }
 tracing = { workspace = true }
 tracing-opentelemetry = { workspace = true }
 trogon-std = { workspace = true }
+uuid = { workspace = true }
 
 [dev-dependencies]
 tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }

rsworkspace/crates/trogon-nats/src/lib.rs

@@ -50,8 +50,8 @@ pub use client::{FlushClient, PublishClient, RequestClient, SubscribeClient};
 pub use connect::{ConnectError, connect};
 pub use messaging::{
     FlushPolicy, NatsError, PublishOperationError, PublishOptions, PublishOptionsBuilder,
-    RetryPolicy, headers_with_trace_context, inject_trace_context, publish, request,
-    request_with_timeout,
+    REQ_ID_HEADER, RetryPolicy, build_request_headers, extract_req_id, headers_with_trace_context,
+    inject_trace_context, publish, request, request_with_timeout,
 };
 
 #[cfg(feature = "test-support")]

rsworkspace/crates/trogon-nats/src/messaging.rs

@@ -23,12 +23,26 @@ pub fn inject_trace_context(headers: &mut HeaderMap) {
     });
 }
 
+pub const REQ_ID_HEADER: &str = "X-Req-Id";
+
 pub fn headers_with_trace_context() -> HeaderMap {
     let mut headers = HeaderMap::new();
     inject_trace_context(&mut headers);
     headers
 }
 
+pub fn build_request_headers() -> HeaderMap {
+    let mut headers = headers_with_trace_context();
+    headers.insert(REQ_ID_HEADER, uuid::Uuid::new_v4().to_string().as_str());
+    headers
+}
+
+pub fn extract_req_id(headers: Option<&HeaderMap>) -> Option<&str> {
+    headers
+        .and_then(|h| h.get(REQ_ID_HEADER))
+        .map(|v| v.as_str())
+}
+
 pub async fn request_with_timeout<N: RequestClient, Req, Res>(
     client: &N,
     subject: &str,
@@ -40,7 +54,7 @@ where
     Res: DeserializeOwned,
 {
     let payload = serde_json::to_vec(request).map_err(NatsError::Serialize)?;
-    let headers = headers_with_trace_context();
+    let headers = build_request_headers();
 
     let response = tokio::time::timeout(
         timeout,