feat(cli): forward TRACEPARENT/TRACESTATE to gateway for CI trace continuity

Cargo.toml

@@ -16,7 +16,7 @@ repository = "https://github.com/NVIDIA/OpenShell"
 # Product backend selection. Consumer crates inherit this choice.
 openshell-crypto = { path = "crates/openshell-crypto", default-features = false, features = ["aws-lc"] }
 # Async runtime
-tokio = { version = "1.43", features = ["full"] }
+tokio = { version = "1.54", features = ["full"] }
 
 # gRPC/Protobuf
 tonic = "0.14"

crates/openshell-cli/Cargo.toml

@@ -18,6 +18,7 @@ path = "src/main.rs"
 openshell-crypto = { workspace = true }
 openshell-bootstrap = { path = "../openshell-bootstrap" }
 openshell-core = { path = "../openshell-core", default-features = false }
+openshell-otel = { path = "../openshell-otel" }
 openshell-policy = { path = "../openshell-policy" }
 openshell-providers = { path = "../openshell-providers" }
 openshell-sdk = { path = "../openshell-sdk" }

crates/openshell-cli/src/tls.rs

@@ -2,7 +2,9 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use miette::{IntoDiagnostic, Result, WrapErr};
-use openshell_sdk::{AuthConfig, ClientConfig, ClientIdentity, OpenShellClient};
+use openshell_otel::EnvTraceContextInterceptor;
+use openshell_sdk::raw::GrpcClient as RawGrpcClient;
+use openshell_sdk::{AuthConfig, ClientConfig, ClientIdentity, EdgeAuthInterceptor};
 use rustls::{
     RootCertStore,
     client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier},
@@ -11,10 +13,34 @@ use rustls::{
 use std::io::Cursor;
 use std::path::PathBuf;
 use tonic::metadata::AsciiMetadataValue;
+use tonic::service::interceptor::InterceptedService;
 use tonic::transport::Channel;
 
 /// Concrete gRPC client type used by all commands.
-pub type GrpcClient = openshell_sdk::raw::AuthedGrpcClient;
+pub type GrpcClient = RawGrpcClient<InterceptedService<Channel, GatewayInterceptor>>;
+
+/// Interceptor stack applied to gateway requests from clients built by
+/// [`grpc_client`].
+///
+/// Runs authentication header injection first, then passive W3C trace-context
+/// forwarding so a CI pipeline's `TRACEPARENT` extends into the gateway's
+/// spans. tonic allows a single interceptor per channel, so the two concerns
+/// are composed here.
+#[derive(Clone)]
+pub struct GatewayInterceptor {
+    auth: EdgeAuthInterceptor,
+    trace: EnvTraceContextInterceptor,
+}
+
+impl tonic::service::Interceptor for GatewayInterceptor {
+    fn call(
+        &mut self,
+        request: tonic::Request<()>,
+    ) -> std::result::Result<tonic::Request<()>, tonic::Status> {
+        let request = self.auth.call(request)?;
+        self.trace.call(request)
+    }
+}
 
 #[derive(Clone, Debug, Default)]
 pub struct TlsOptions {
@@ -303,19 +329,30 @@ pub async fn build_channel(server: &str, tls: &TlsOptions) -> Result<Channel> {
         .into_diagnostic()
 }
 
-/// Build the SDK client and return its raw gRPC escape hatch.
+/// Build the gRPC client used by all commands.
 ///
-/// When `tls.edge_token` is set, the returned client is wrapped with an
-/// interceptor that injects authentication headers on every request.
-/// Otherwise, standard mTLS is used (interceptor is a no-op).
+/// The client is built on the SDK's channel construction (TLS discovery,
+/// edge tunnel, insecure mode) and carries a [`GatewayInterceptor`] that
+/// injects authentication headers (when a token is set) and passively
+/// forwards the `TRACEPARENT`/`TRACESTATE` trace context from the
+/// environment on every request. With no token and no trace context, the
+/// interceptor is a no-op.
 pub async fn grpc_client(server: &str, tls: &TlsOptions) -> Result<GrpcClient> {
-    let config = sdk_client_config(server, tls)?;
-    // Keep the SDK's transport state machine out of every command future.
-    // Most CLI commands compose this helper into larger async operations.
-    let client = Box::pin(OpenShellClient::connect(config))
-        .await
+    // Keep the SDK's channel-building state machine out of every command
+    // future. Most CLI commands compose this helper into larger async
+    // operations.
+    let channel = Box::pin(build_channel(server, tls)).await?;
+    let interceptor = interceptor_from_tls(tls)?;
+    Ok(RawGrpcClient::with_interceptor(channel, interceptor))
+}
+
+fn interceptor_from_tls(tls: &TlsOptions) -> Result<GatewayInterceptor> {
+    let auth = EdgeAuthInterceptor::new(tls.oidc_token.as_deref(), tls.edge_token.as_deref())
         .into_diagnostic()?;
-    Ok(client.raw_grpc())
+    Ok(GatewayInterceptor {
+        auth,
+        trace: EnvTraceContextInterceptor::from_env(),
+    })
 }
 
 /// Gateway replica that owns the sandbox named in `response`, when it says.
@@ -365,13 +402,61 @@ pub fn retry_unrouted(replica: Option<&AsciiMetadataValue>, status: &tonic::Stat
 #[cfg(test)]
 mod tests {
     use super::{
-        TlsOptions, owner_replica, retry_unrouted, routed_to, sdk_client_config, with_owner_hint,
+        TlsOptions, interceptor_from_tls, owner_replica, retry_unrouted, routed_to,
+        sdk_client_config, with_owner_hint,
     };
+    use crate::TEST_ENV_LOCK;
     use openshell_core::replica_routing::{
         OWNER_REPLICA_HEADER, OWNER_REQUEST_HEADER, ROUTE_REPLICA_HEADER,
     };
     use openshell_sdk::AuthConfig;
     use std::io::Write;
+    use tonic::service::Interceptor as _;
+
+    #[test]
+    fn gateway_interceptor_composes_auth_and_trace_forwarding() {
+        let _guard = TEST_ENV_LOCK
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        let traceparent = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";
+        temp_env::with_vars(
+            [
+                ("TRACEPARENT", Some(traceparent)),
+                ("TRACESTATE", Some("vendor=value")),
+            ],
+            || {
+                let tls = TlsOptions {
+                    oidc_token: Some("test-token".to_string()),
+                    ..TlsOptions::new(None, None, None)
+                };
+                let mut interceptor = interceptor_from_tls(&tls).unwrap();
+
+                let request = interceptor.call(tonic::Request::new(())).unwrap();
+
+                assert_eq!(
+                    request
+                        .metadata()
+                        .get("authorization")
+                        .and_then(|value| value.to_str().ok()),
+                    Some("Bearer test-token")
+                );
+                assert_eq!(
+                    request
+                        .metadata()
+                        .get("traceparent")
+                        .and_then(|value| value.to_str().ok()),
+                    Some(traceparent)
+                );
+                assert_eq!(
+                    request
+                        .metadata()
+                        .get("tracestate")
+                        .and_then(|value| value.to_str().ok()),
+                    Some("vendor=value")
+                );
+            },
+        );
+    }
 
     #[test]
     fn owner_hint_is_requested_explicitly() {

crates/openshell-core/Cargo.toml

@@ -46,6 +46,7 @@ openshell-otel = { path = "../openshell-otel", optional = true }
 [target.'cfg(unix)'.dependencies]
 nix = { workspace = true }
 rustix = { workspace = true }
+socket2 = { workspace = true }
 
 [features]
 default = ["telemetry"]

crates/openshell-core/src/net.rs

@@ -14,6 +14,10 @@
 //! [`connect_tcp_nodelay_best_effort`]) help avoid the known latency
 //! imposed by conflict between Nagle's algorithm and delayed ACK behaviors.
 //!
+//! The `AsyncFd` registration helpers ([`register_socket_async_fd`],
+//! [`register_file_async_fd`]) wrap tokio's I/O-safety-`unsafe`
+//! `AsyncFd::register` for owned descriptor types that uphold its contract.
+//!
 
 use ipnet::{IpNet, Ipv4Net, Ipv6Net};
 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
@@ -285,6 +289,40 @@ fn is_internal_v4(v4: Ipv4Addr) -> bool {
     false
 }
 
+/// Register an owned [`socket2::Socket`] with the current tokio reactor.
+///
+/// Tokio's `AsyncFd::register` is `unsafe` because a general `AsRawFd`
+/// implementation can expose a descriptor that is closed, reused, or shared
+/// while the `AsyncFd` still holds it. A `socket2::Socket` owns its descriptor
+/// exclusively and keeps it open and referring to the same file description
+/// until dropped, which upholds the contract, so this wrapper is safe.
+#[cfg(unix)]
+#[allow(unsafe_code)]
+pub fn register_socket_async_fd(
+    socket: socket2::Socket,
+) -> std::io::Result<tokio::io::unix::AsyncFd<socket2::Socket>> {
+    // SAFETY: `socket` is moved into the returned `AsyncFd`, and a
+    // `socket2::Socket` keeps its descriptor valid, open, and referring to the
+    // same file description for the value's entire lifetime.
+    unsafe { tokio::io::unix::AsyncFd::register(socket) }.map_err(std::io::Error::from)
+}
+
+/// Register an owned [`std::fs::File`] with the current tokio reactor.
+///
+/// Sound for the same reason as [`register_socket_async_fd`]: a `File` owns
+/// its descriptor exclusively and keeps it open and referring to the same
+/// file description until dropped.
+#[cfg(unix)]
+#[allow(unsafe_code)]
+pub fn register_file_async_fd(
+    file: std::fs::File,
+) -> std::io::Result<tokio::io::unix::AsyncFd<std::fs::File>> {
+    // SAFETY: `file` is moved into the returned `AsyncFd`, and a `File`
+    // keeps its descriptor valid, open, and referring to the same file
+    // description for the value's entire lifetime.
+    unsafe { tokio::io::unix::AsyncFd::register(file) }.map_err(std::io::Error::from)
+}
+
 /// Enable `TCP_NODELAY` on a stream, logging (not returning) any failure.
 ///
 /// Disabling Nagle's algorithm keeps small writes from waiting on delayed ACKs.

crates/openshell-otel/src/lib.rs

@@ -21,9 +21,9 @@ pub use grpc::{
 };
 pub use ocsf::{OcsfOtlp, OcsfOtlpLayer};
 pub use propagation::{
-    HeaderMapExtractor, MetadataMapInjector, TRACEPARENT_ENV, TRACESTATE_ENV,
-    TraceContextInterceptor, current_trace_context_carrier, current_trace_context_environment,
-    set_parent_from_environment, trace_context_environment,
+    EnvTraceContextInterceptor, HeaderMapExtractor, MetadataMapInjector, TRACEPARENT_ENV,
+    TRACESTATE_ENV, TraceContextInterceptor, current_trace_context_carrier,
+    current_trace_context_environment, set_parent_from_environment, trace_context_environment,
 };
 
 use opentelemetry::KeyValue;