Utility-Protocol · elizabetheonoja-art · Aug 19, 2026 · Aug 19, 2026
diff --git a/src/api/middleware.rs b/src/api/middleware.rs
@@ -12,7 +12,7 @@ use std::net::SocketAddr;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 use std::time::{Duration, Instant};
-use tracing::warn;
+use tracing::{warn, Instrument};
 
 lazy_static::lazy_static! {
     static ref START_INSTANT: Instant = Instant::now();
@@ -398,6 +398,23 @@ pub async fn rate_limit_layer(
     next.run(req).await
 }
 
+pub async fn correlation_id_layer(mut req: Request<Body>, next: Next) -> Response {
+    let correlation_id = crate::tracing::correlation::correlation_id_from_headers(req.headers());
+    req.extensions_mut().insert(correlation_id.clone());
+
+    let span = tracing::info_span!(
+        "http.correlation",
+        correlation_id = %correlation_id.as_str()
+    );
+
+    let mut response = crate::tracing::correlation::scope(correlation_id.clone(), async move {
+        next.run(req).instrument(span).await
+    })
+    .await;
+    crate::tracing::correlation::insert_header(response.headers_mut(), &correlation_id);
+    response
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

diff --git a/src/api/router.rs b/src/api/router.rs
@@ -109,6 +109,9 @@ pub async fn build_router(state: AppState) -> anyhow::Result<Router> {
         .layer(axum_mw::from_fn(
             crate::api::middleware::slo_monitoring_layer,
         ))
+        .layer(axum_mw::from_fn(
+            crate::api::middleware::correlation_id_layer,
+        ))
         .layer(cors)
         .with_state(state);
 

diff --git a/src/service_mesh/client.rs b/src/service_mesh/client.rs
@@ -144,9 +144,12 @@ impl ServiceMeshClient {
             req = req.header("content-type", "application/json").body(b);
         }
 
-        req.header("x-service-name", &self.identity.service_account)
+        req = req
+            .header("x-service-name", &self.identity.service_account)
             .header("x-trust-domain", &self.identity.trust_domain)
-            .timeout(std::time::Duration::from_secs(endpoint.timeout_secs))
+            .timeout(std::time::Duration::from_secs(endpoint.timeout_secs));
+
+        crate::tracing::correlation::propagate_reqwest(req)
             .send()
             .await
             .map_err(|e| ServiceMeshClientError::HttpError(e.to_string()))
@@ -164,10 +167,12 @@ impl ServiceMeshClient {
             .join(endpoint.health_check_path.trim_start_matches('/'))
             .map_err(|e| ServiceMeshClientError::HttpError(e.to_string()))?;
 
-        match self
+        let req = self
             .http_client
             .get(url)
-            .timeout(std::time::Duration::from_secs(5))
+            .timeout(std::time::Duration::from_secs(5));
+
+        match crate::tracing::correlation::propagate_reqwest(req)
             .send()
             .await
         {

diff --git a/src/tracing/correlation.rs b/src/tracing/correlation.rs
@@ -0,0 +1,143 @@
+use std::collections::HashMap;
+use std::future::Future;
+
+use axum::http::{HeaderMap, HeaderValue};
+use uuid::Uuid;
+
+pub const CORRELATION_ID_HEADER: &str = "x-correlation-id";
+const MAX_CORRELATION_ID_LEN: usize = 128;
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct CorrelationId(String);
+
+impl CorrelationId {
+    pub fn generate() -> Self {
+        Self(Uuid::new_v4().to_string())
+    }
+
+    pub fn parse(value: &str) -> Option<Self> {
+        let trimmed = value.trim();
+        if trimmed.is_empty() || trimmed.len() > MAX_CORRELATION_ID_LEN {
+            return None;
+        }
+
+        if trimmed
+            .chars()
+            .all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '-' | '_' | '.'))
+        {
+            Some(Self(trimmed.to_string()))
+        } else {
+            None
+        }
+    }
+
+    pub fn as_str(&self) -> &str {
+        &self.0
+    }
+
+    fn to_header_value(&self) -> HeaderValue {
+        HeaderValue::from_str(self.as_str()).expect("correlation ids are sanitized")
+    }
+}
+
+tokio::task_local! {
+    static CURRENT_CORRELATION_ID: String;
+}
+
+pub async fn scope<F, T>(id: CorrelationId, future: F) -> T
+where
+    F: Future<Output = T>,
+{
+    CURRENT_CORRELATION_ID.scope(id.0, future).await
+}
+
+pub fn current_correlation_id() -> Option<String> {
+    CURRENT_CORRELATION_ID.try_with(Clone::clone).ok()
+}
+
+pub fn correlation_id_from_headers(headers: &HeaderMap) -> CorrelationId {
+    headers
+        .get(CORRELATION_ID_HEADER)
+        .and_then(|value| value.to_str().ok())
+        .and_then(CorrelationId::parse)
+        .unwrap_or_else(CorrelationId::generate)
+}
+
+pub fn insert_header(headers: &mut HeaderMap, id: &CorrelationId) {
+    headers.insert(CORRELATION_ID_HEADER, id.to_header_value());
+}
+
+pub fn inject_current_into_message_headers(headers: &mut HashMap<String, String>) {
+    if let Some(id) = current_correlation_id() {
+        headers.insert(CORRELATION_ID_HEADER.to_string(), id);
+    }
+}
+
+pub fn inject_into_message_headers(headers: &mut HashMap<String, String>, id: &CorrelationId) {
+    headers.insert(CORRELATION_ID_HEADER.to_string(), id.as_str().to_string());
+}
+
+pub fn extract_from_message_headers(headers: &HashMap<String, String>) -> Option<CorrelationId> {
+    headers
+        .get(CORRELATION_ID_HEADER)
+        .or_else(|| headers.get("X-Correlation-ID"))
+        .and_then(|value| CorrelationId::parse(value))
+}
+
+pub fn propagate_reqwest(builder: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
+    match current_correlation_id() {
+        Some(id) => builder.header(CORRELATION_ID_HEADER, id),
+        None => builder,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn parses_only_safe_correlation_ids() {
+        assert_eq!(
+            CorrelationId::parse("abc-123_DEF.456").unwrap().as_str(),
+            "abc-123_DEF.456"
+        );
+        assert!(CorrelationId::parse("").is_none());
+        assert!(CorrelationId::parse("has spaces").is_none());
+        assert!(CorrelationId::parse("secret\nnext").is_none());
+        assert!(CorrelationId::parse(&"a".repeat(MAX_CORRELATION_ID_LEN + 1)).is_none());
+    }
+
+    #[test]
+    fn extracts_existing_header_or_generates_new_id() {
+        let mut headers = HeaderMap::new();
+        headers.insert(CORRELATION_ID_HEADER, HeaderValue::from_static("incoming-1"));
+        assert_eq!(correlation_id_from_headers(&headers).as_str(), "incoming-1");
+
+        headers.insert(CORRELATION_ID_HEADER, HeaderValue::from_static("bad id"));
+        let generated = correlation_id_from_headers(&headers);
+        assert_ne!(generated.as_str(), "bad id");
+        assert!(Uuid::parse_str(generated.as_str()).is_ok());
+    }
+
+    #[tokio::test]
+    async fn task_scope_exposes_current_correlation_id() {
+        let id = CorrelationId::parse("req-123").unwrap();
+        let current = scope(id, async { current_correlation_id() }).await;
+        assert_eq!(current.as_deref(), Some("req-123"));
+        assert!(current_correlation_id().is_none());
+    }
+
+    #[tokio::test]
+    async fn message_headers_roundtrip_current_id() {
+        let id = CorrelationId::parse("event-456").unwrap();
+        let mut headers = HashMap::new();
+        scope(id, async {
+            inject_current_into_message_headers(&mut headers);
+        })
+        .await;
+        assert_eq!(
+            extract_from_message_headers(&headers).unwrap().as_str(),
+            "event-456"
+        );
+    }
+}
diff --git a/src/tracing/kafka_propagator.rs b/src/tracing/kafka_propagator.rs
@@ -31,6 +31,7 @@ pub fn inject_into_kafka_headers(cx: &Context) -> HashMap<String, String> {
     global::get_text_map_propagator(|propagator| {
         propagator.inject_context(cx, &mut KafkaHeaderInjector(&mut headers));
     });
+    crate::tracing::correlation::inject_current_into_message_headers(&mut headers);
     headers
 }
 

diff --git a/src/tracing/logging.rs b/src/tracing/logging.rs
@@ -198,6 +198,12 @@ where
                 .unwrap()
                 .insert("span_id".to_string(), serde_json::Value::String(s_id));
         }
+        if let Some(correlation_id) = crate::tracing::correlation::current_correlation_id() {
+            log_record.as_object_mut().unwrap().insert(
+                "correlation_id".to_string(),
+                serde_json::Value::String(correlation_id),
+            );
+        }
 
         if !visitor.attributes.is_empty() {
             log_record.as_object_mut().unwrap().insert(

diff --git a/src/tracing/mod.rs b/src/tracing/mod.rs
@@ -1,3 +1,4 @@
+pub mod correlation;
 pub mod db_tracing;
 pub mod exporters;
 pub mod kafka_propagator;

diff --git a/tests/distributed_tracing_test.rs b/tests/distributed_tracing_test.rs
@@ -17,6 +17,9 @@ use utility_backend::gateway::telemetry::{
 use utility_backend::tracing::kafka_propagator::{
     extract_from_kafka_headers, inject_into_kafka_headers,
 };
+use utility_backend::tracing::correlation::{
+    extract_from_message_headers, scope, CorrelationId, CORRELATION_ID_HEADER,
+};
 
 // ── OTLP Exporter & TracerProvider ──────────────────────────────────
 
@@ -114,6 +117,24 @@ async fn test_kafka_headers_inject_extract_roundtrip() {
     );
 }
 
+#[tokio::test]
+async fn test_kafka_headers_include_current_correlation_id() {
+    let _ = init_open_telemetry("kafka-correlation-test");
+    let id = CorrelationId::parse("corr-789").unwrap();
+    let parent = opentelemetry::Context::new();
+
+    let headers = scope(id, async { inject_into_kafka_headers(&parent) }).await;
+
+    assert_eq!(
+        headers.get(CORRELATION_ID_HEADER).map(String::as_str),
+        Some("corr-789")
+    );
+    assert_eq!(
+        extract_from_message_headers(&headers).unwrap().as_str(),
+        "corr-789"
+    );
+}
+
 #[tokio::test]
 async fn test_kafka_headers_empty_input_produces_empty_context() {
     let _ = init_open_telemetry("kafka-empty-test");

diff --git a/tests/logging_test.rs b/tests/logging_test.rs
@@ -5,6 +5,7 @@ use std::io;
 use std::sync::{Arc, Mutex};
 use tracing::{info, span, Level};
 use tracing_subscriber::{filter::LevelFilter, layer::SubscriberExt, Registry};
+use utility_backend::tracing::correlation::{scope, CorrelationId};
 use utility_backend::tracing::logging::OtelJsonFormatter;
 
 #[derive(Clone)]
@@ -62,13 +63,17 @@ async fn test_otel_json_logging_format() {
         .with(otel_layer)
         .with(fmt_layer);
 
-    // Run within our custom subscriber
-    tracing::subscriber::with_default(subscriber, || {
-        let test_span = span!(Level::INFO, "test_operation");
-        let _enter = test_span.enter();
+    let correlation_id = CorrelationId::parse("log-correlation-1").unwrap();
+    scope(correlation_id, async {
+        // Run within our custom subscriber
+        tracing::subscriber::with_default(subscriber, || {
+            let test_span = span!(Level::INFO, "test_operation");
+            let _enter = test_span.enter();
 
-        info!("structured log inside span");
-    });
+            info!("structured log inside span");
+        });
+    })
+    .await;
 
     let output_bytes = buffer.lock().unwrap().clone();
     let output_str = String::from_utf8(output_bytes).unwrap();
@@ -90,4 +95,5 @@ async fn test_otel_json_logging_format() {
     assert!(log2["span_id"].is_string(), "Missing span_id in span log");
     assert!(!log2["trace_id"].as_str().unwrap().is_empty());
     assert!(!log2["span_id"].as_str().unwrap().is_empty());
+    assert_eq!(log2["correlation_id"], "log-correlation-1");
 }
diff --git a/tests/rate_limit_integration.rs b/tests/rate_limit_integration.rs
@@ -1,5 +1,5 @@
 use axum::{
-    body::Body,
+    body::{to_bytes, Body},
     extract::ConnectInfo,
     http::{Request, StatusCode},
     middleware as axum_mw,
@@ -8,7 +8,11 @@ use axum::{
 };
 use std::net::SocketAddr;
 use utility_backend::api::middleware::{
-    rate_limit_layer, tenant_rate_limit_layer, DynamicRateLimiter, TenantRateLimiter,
+    correlation_id_layer, rate_limit_layer, tenant_rate_limit_layer, DynamicRateLimiter,
+    TenantRateLimiter,
+};
+use utility_backend::tracing::correlation::{
+    current_correlation_id, CORRELATION_ID_HEADER,
 };
 
 #[tokio::test]
@@ -113,3 +117,31 @@ async fn test_tenant_rate_limit_anonymous() {
     let res = send_no_header(app.clone()).await;
     assert_eq!(res.status(), StatusCode::TOO_MANY_REQUESTS);
 }
+
+#[tokio::test]
+async fn test_correlation_id_middleware_propagates_context_and_response_header() {
+    let app = Router::new()
+        .route(
+            "/",
+            get(|| async { current_correlation_id().unwrap_or_else(|| "missing".to_string()) }),
+        )
+        .layer(axum_mw::from_fn(correlation_id_layer));
+
+    let req = Request::builder()
+        .uri("/")
+        .header(CORRELATION_ID_HEADER, "request-abc")
+        .body(Body::empty())
+        .unwrap();
+
+    let res = tower::ServiceExt::oneshot(app, req).await.unwrap();
+    assert_eq!(res.status(), StatusCode::OK);
+    assert_eq!(
+        res.headers()
+            .get(CORRELATION_ID_HEADER)
+            .and_then(|value| value.to_str().ok()),
+        Some("request-abc")
+    );
+
+    let body = to_bytes(res.into_body(), usize::MAX).await.unwrap();
+    assert_eq!(body.as_ref(), b"request-abc");
+}
-Original file line number
+Diff line change
@@ Expand Up @@
         global::get_text_map_propagator(|propagator| {
             propagator.inject_context(cx, &mut KafkaHeaderInjector(&mut headers));
         });
+        crate::tracing::correlation::inject_current_into_message_headers(&mut headers);
         headers
     }
@@ Expand Down @@