From 587f6bf8a897b522501b69c7d45bf1eaf0f9ce29 Mon Sep 17 00:00:00 2001
From: Floriane TUERNAL SABOTINOV <ftuertinov@proton.me>
Date: Mon, 8 Sep 2025 10:18:28 +0200
Subject: [PATCH] add method to extract trace_id from ContextCarrier

---
 .../bin-common/src/opentelemetry/context.rs   | 12 +++++++
 .../websocket/connection_handler/fresh.rs     | 33 +++++++++++++------
 2 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/common/bin-common/src/opentelemetry/context.rs b/common/bin-common/src/opentelemetry/context.rs
index 76ebb51f7c..35f3357c01 100644
--- a/common/bin-common/src/opentelemetry/context.rs
+++ b/common/bin-common/src/opentelemetry/context.rs
@@ -1,4 +1,5 @@
 use opentelemetry::propagation::{Injector, Extractor};
+use opentelemetry::trace::TraceId;
 use std::collections::HashMap;
 
 /// Make a Carrier for context propagation
@@ -20,6 +21,17 @@ impl ContextCarrier {
     pub fn into_map(self) -> HashMap<String, String> {
         self.data
     }
+
+    pub fn extract_trace_id(&self) -> Option<TraceId> {
+        self.get("traceparent").and_then(|tp| {
+            let parts: Vec<&str> = tp.split('-').collect();
+            if parts.len() == 4 {
+                TraceId::from_hex(parts[1]).ok()
+            } else {
+                None
+            }
+        })
+    }
 }
 
 impl Injector for ContextCarrier {
diff --git a/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs b/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs
index 6e8de34398..f42e007ea6 100644
--- a/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs
+++ b/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs
@@ -35,6 +35,7 @@ use nym_sphinx::DestinationAddressBytes;
 use nym_task::ShutdownToken;
 use opentelemetry::trace::TraceContextExt;
 use opentelemetry::TraceId;
+use opentelemetry::Context;
 use opentelemetry_sdk::logs::TraceContext;
 use opentelemetry_sdk::trace::{IdGenerator, RandomIdGenerator};
 use rand::CryptoRng;
@@ -869,21 +870,38 @@ impl<R, S> FreshHandler<R, S> {
         use opentelemetry_sdk::propagation::TraceContextPropagator;
         use opentelemetry::propagation::TextMapPropagator;
 
-        if let ClientControlRequest::AuthenticateV2(ref auth_req) = request {
+        let remote_cx_span = if let ClientControlRequest::AuthenticateV2(ref auth_req) = request {
             if let Some(otel_context) = &auth_req.otel_context {
                 let carrier = ContextCarrier::from_map(otel_context.clone());
+
+                let extracted_trace_id =  carrier.extract_trace_id();
+                info!("Extracted trace id: {:?}", extracted_trace_id);
+
                 let propagator = TraceContextPropagator::new();
                 let extracted_context = propagator.extract(&carrier);
-                let _guard = extracted_context.attach();
 
+
+                let span = info_span!("extracted_otel_context");
+                tracing::Span::current().set_parent(extracted_context);
                 error!("==== Context propagation successful ====");
+                
+                Some(span)
             } else {
                 warn!("No OpenTelemetry context provided in the request");
+                None
             }
-        }
+        } else {
+            None
+        };
 
-        let span = info_span!("websocket_authentication", request = ?request);
-        let _ = span.enter();
+        let child_span = if let Some(ref parent_span) = remote_cx_span {
+            info_span!(parent: parent_span, "handling_initial_client_request")
+        } else {
+            info_span!("handling_initial_client_request")
+        };
+        let _enter = child_span.enter();
+
+        tracing::error!("handling initial client request");
 
         let auth_result = match request {
             ClientControlRequest::Authenticate {
@@ -974,11 +992,6 @@ impl<R, S> FreshHandler<R, S> {
                 }
             };
 
-                info_span!("THE SOLE PURPOSE OF THIS SPAN IS TO SEE ITS TRACE ID");
-                let current_context = opentelemetry::Context::current();
-                let final_trace_id = current_context.span().span_context().trace_id();
-                error!("==== trace_id after initial client request handling returns: {:?} ====", final_trace_id);
-
             if let Some(registration_details) = maybe_auth_res {
                 let (mix_sender, mix_receiver) = mpsc::unbounded();
                 // Channel for handlers to ask other handlers if they are still active.