Publish harness and TUI open-source

initial sync from the monorepo

crates/common/xai-tracing/Cargo.toml

@@ -0,0 +1,37 @@
+[package]
+license = "Apache-2.0"
+edition.workspace = true
+name = "xai-tracing"
+version = "0.1.0"
+
+[dependencies]
+async-trait = { workspace = true }
+http = { workspace = true }
+log = { workspace = true }
+fastrace = { workspace = true }
+fastrace-tonic = { workspace = true }
+fastrace-opentelemetry = { workspace = true }
+fastrace-reqwest = { workspace = true }
+opentelemetry = { workspace = true }
+opentelemetry-http = { workspace = true }
+opentelemetry-otlp = { workspace = true }
+opentelemetry_sdk = { workspace = true, features = ["testing"] }
+reqwest = { workspace = true }
+reqwest-middleware = { workspace = true }
+tokio = { workspace = true }
+tonic = { workspace = true }
+tower = { workspace = true, features = ["full"] }
+tower-http = { workspace = true, features = ["trace"] }
+tracing = { workspace = true }
+tracing-opentelemetry = { workspace = true }
+uuid = { workspace = true }
+
+[dev-dependencies]
+bytes = { workspace = true }
+http-body-util = { workspace = true }
+tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
+tracing-subscriber = { workspace = true }
+wiremock = { workspace = true }
+
+[lints]
+workspace = true

crates/common/xai-tracing/src/dispatch.rs

@@ -0,0 +1,45 @@
+use tracing::subscriber::NoSubscriber;
+
+/// Returns `true` when a `tracing` dispatcher (subscriber) is active in the
+/// current context — either the thread-scoped default
+/// (`tracing::subscriber::with_default` / `set_default`) or the global one.
+///
+/// When this returns `false`, spans have no consumer. Worse than useless: when
+/// `tracing` is compiled with its `log` compatibility feature, every span
+/// creation and every later `Span::record(...)` is downgraded to a `log`
+/// record at the span's level. In processes that only configure a `log`
+/// logger — e.g. integration tests or fastrace-only binaries — that prints
+/// noise like:
+///
+/// ```text
+/// I grpc; otel.name="POST_/pkg.Service/Method" ...
+/// I grpc; status_code=200 OK
+/// I grpc; trace_id=00000000000000000000000000000000
+/// ```
+///
+/// Request-span factories (gRPC/HTTP server and client middleware) call this
+/// and return `Span::none()` when no dispatcher is active, so the span is
+/// neither built nor downgraded to log spam.
+pub fn dispatcher_active() -> bool {
+    tracing::dispatcher::get_default(|dispatch| !dispatch.is::<NoSubscriber>())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // NOTE: relies on no test in this binary installing a *global* subscriber
+    // (`OtelTestEnv` and friends use thread-scoped `set_default` guards).
+    #[test]
+    fn without_dispatcher_inactive() {
+        assert!(!dispatcher_active());
+    }
+
+    #[test]
+    fn scoped_dispatcher_active() {
+        tracing::subscriber::with_default(tracing_subscriber::registry(), || {
+            assert!(dispatcher_active());
+        });
+        assert!(!dispatcher_active());
+    }
+}

crates/common/xai-tracing/src/fastrace.rs

@@ -0,0 +1,83 @@
+use fastrace::prelude::*;
+use fastrace_opentelemetry::OpenTelemetryReporter;
+use opentelemetry::InstrumentationScope;
+use opentelemetry::KeyValue;
+use opentelemetry_otlp::WithExportConfig;
+use opentelemetry_otlp::{ExporterBuildError, SpanExporter};
+use opentelemetry_sdk::Resource;
+use std::borrow::Cow;
+use std::iter;
+
+// Fastrace initialization
+pub fn init_fastrace(
+    endpoint: String,
+    name: String,
+    resource_attributes: impl IntoIterator<Item = (String, String)>,
+) -> Result<(), ExporterBuildError> {
+    let exporter = SpanExporter::builder()
+        .with_tonic()
+        .with_endpoint(endpoint)
+        .with_protocol(opentelemetry_otlp::Protocol::Grpc)
+        .with_timeout(opentelemetry_otlp::OTEL_EXPORTER_OTLP_TIMEOUT_DEFAULT)
+        .build()?;
+    let attributes = resource_attributes
+        .into_iter()
+        .chain(iter::once(("service.name".into(), name.clone())))
+        .map(|(k, v)| KeyValue::new(k, v));
+    let reporter = OpenTelemetryReporter::new(
+        exporter,
+        Cow::Owned(Resource::builder().with_attributes(attributes).build()),
+        InstrumentationScope::builder(name)
+            .with_version(env!("CARGO_PKG_VERSION"))
+            .build(),
+    );
+    fastrace::set_reporter(reporter, fastrace::collector::Config::default());
+    Ok(())
+}
+
+pub fn current_trace_id() -> Option<String> {
+    SpanContext::current_local_parent().map(|current| current.encode_w3c_traceparent())
+}
+
+pub fn local_or_random_span_ctx() -> SpanContext {
+    SpanContext::current_local_parent().unwrap_or_else(SpanContext::random)
+}
+
+pub fn enter_span_with_traceparent(name: impl Into<Cow<'static, str>>, traceparent: &str) -> Span {
+    if let Some(span_ctx) = SpanContext::decode_w3c_traceparent(traceparent) {
+        Span::root(name, span_ctx)
+    } else {
+        Span::enter_with_local_parent(name)
+    }
+}
+
+// Tonic channel (TODO: Move into grpc_client when deprecated tracing)
+#[allow(dead_code)]
+pub type FastraceChannel = fastrace_tonic::FastraceClientService<tonic::transport::Channel>;
+
+pub fn fastrace_channel(
+    channel: tonic::transport::Channel,
+) -> fastrace_tonic::FastraceClientService<tonic::transport::Channel> {
+    tower::ServiceBuilder::new()
+        .layer(fastrace_tonic::FastraceClientLayer)
+        .service(channel)
+}
+
+// Request middleware (TODO: Move into http_client when deprecated tracing)
+#[derive(Clone)]
+#[allow(dead_code)]
+pub struct TraceparentMiddleware;
+
+#[async_trait::async_trait]
+impl reqwest_middleware::Middleware for TraceparentMiddleware {
+    async fn handle(
+        &self,
+        mut req: reqwest::Request,
+        extensions: &mut http::Extensions,
+        next: reqwest_middleware::Next<'_>,
+    ) -> reqwest_middleware::Result<reqwest::Response> {
+        req.headers_mut()
+            .extend(fastrace_reqwest::traceparent_headers());
+        next.run(req, extensions).await
+    }
+}

crates/common/xai-tracing/src/grpc_client.rs

@@ -0,0 +1,388 @@
+use http::{HeaderMap, Request};
+use opentelemetry::{global, propagation::Extractor, propagation::Injector};
+use std::task::{Context, Poll};
+use tonic::transport::Channel;
+use tonic::{
+    Status,
+    metadata::{MetadataKey, MetadataMap, MetadataValue},
+};
+use tower::{Layer, Service, ServiceBuilder};
+use tower_http::classify::{GrpcErrorsAsFailures, SharedClassifier};
+use tower_http::trace::{MakeSpan, Trace, TraceLayer};
+use tracing::{Span, warn};
+use tracing_opentelemetry::OpenTelemetrySpanExt;
+
+pub type TracedChannel = Trace<
+    InjectTraceContextService<Channel>,
+    SharedClassifier<GrpcErrorsAsFailures>,
+    MakeClientSpan,
+>;
+
+/// Wraps the input channel with a tracing layer. This function can be used to create a traced gRPC
+/// client as follows:
+///
+/// ```rust
+/// use tonic::transport::Endpoint;
+/// use std::str::FromStr;
+/// use xai_tracing::traced_channel;
+///
+/// let channel = Endpoint::from_str("http://foo").unwrap();
+/// //let client = SomeClient::new(traced_channel(channel));
+///```
+pub fn traced_channel(channel: Channel) -> TracedChannel {
+    ServiceBuilder::new()
+        .layer(TraceLayer::new_for_grpc().make_span_with(MakeClientSpan))
+        .layer(InjectTraceContextLayer)
+        .service(channel)
+}
+
+/// Implements the [`MakeSpan`] trait, to trace outgoing gRPC requests.
+#[derive(Debug, Clone, Copy)]
+pub struct MakeClientSpan;
+
+impl<B> MakeSpan<B> for MakeClientSpan {
+    fn make_span(&mut self, request: &Request<B>) -> Span {
+        // No active dispatcher → the span has no consumer and would only be
+        // downgraded to `log` spam. See `crate::dispatcher_active`.
+        if !crate::dispatcher_active() {
+            return Span::none();
+        }
+        tracing::info_span!(
+            "grpc_request",
+            otel.kind = "client",
+            method = %request.method(),
+            uri = %request.uri(),
+            version = ?request.version(),
+        )
+    }
+}
+
+#[derive(Clone, Copy, Debug, Default)]
+pub struct InjectTraceContextLayer;
+
+impl<S> Layer<S> for InjectTraceContextLayer {
+    type Service = InjectTraceContextService<S>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        InjectTraceContextService { inner }
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct InjectTraceContextService<S> {
+    inner: S,
+}
+
+impl<S, B> Service<Request<B>> for InjectTraceContextService<S>
+where
+    S: Service<Request<B>>,
+{
+    type Response = S::Response;
+    type Error = S::Error;
+    type Future = S::Future;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, mut req: Request<B>) -> Self::Future {
+        crate::http_client::attach_trace_to_http_request(req.headers_mut());
+        self.inner.call(req)
+    }
+}
+
+/// Inject W3C `traceparent` / `tracestate` from the active span into gRPC
+/// metadata. Mutates in place so callers never lose the request body.
+pub fn attach_trace_to_grpc_request_mut(metadata: &mut MetadataMap) {
+    global::get_text_map_propagator(|propagator| {
+        let context = Span::current().context();
+        propagator.inject_context(&context, &mut MetadataInjector(metadata));
+    });
+}
+
+/// Trace context propagation: send the trace context by injecting it into the metadata of the given
+/// request.
+pub fn attach_trace_to_grpc_request<T>(
+    mut request: tonic::Request<T>,
+) -> Result<tonic::Request<T>, Status> {
+    attach_trace_to_grpc_request_mut(request.metadata_mut());
+    Ok(request)
+}
+
+// Need a custom Injector to inject OTel headers
+pub struct MetadataInjector<'a>(&'a mut MetadataMap);
+
+impl Injector for MetadataInjector<'_> {
+    fn set(&mut self, key: &str, value: String) {
+        match MetadataKey::from_bytes(key.as_bytes()) {
+            Ok(key) => match MetadataValue::try_from(&value) {
+                Ok(value) => {
+                    self.0.insert(key, value);
+                }
+
+                Err(error) => warn!(value, error = format!("{error:#}"), "parse metadata value"),
+            },
+
+            Err(error) => warn!(key, error = format!("{error:#}"), "parse metadata key"),
+        }
+    }
+}
+
+pub struct HeaderExtractor<'a>(pub &'a HeaderMap);
+
+impl Extractor for HeaderExtractor<'_> {
+    fn get(&self, key: &str) -> Option<&str> {
+        self.0.get(key).and_then(|v| {
+            let s = v.to_str();
+            if let Err(ref error) = s {
+                warn!(%error, ?v, "cannot convert header value to ASCII")
+            };
+            s.ok()
+        })
+    }
+
+    fn keys(&self) -> Vec<&str> {
+        self.0.keys().map(|k| k.as_str()).collect()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::testing::{OtelTestEnv, otel_span_id_hex, otel_trace_id_hex, parse_traceparent};
+    use http_body_util::Empty;
+    use std::convert::Infallible;
+    use std::sync::{Arc, Mutex};
+    use std::time::Duration;
+    use tower_http::classify::GrpcFailureClass;
+    use tracing::Instrument;
+
+    type EmptyBody = Empty<bytes::Bytes>;
+
+    #[derive(Clone)]
+    struct CaptureService {
+        seen: Arc<Mutex<Option<HeaderMap>>>,
+        response_grpc_status: Option<&'static str>,
+    }
+
+    impl CaptureService {
+        fn new() -> Self {
+            Self {
+                seen: Arc::new(Mutex::new(None)),
+                response_grpc_status: None,
+            }
+        }
+
+        fn with_grpc_status(status: &'static str) -> Self {
+            Self {
+                seen: Arc::new(Mutex::new(None)),
+                response_grpc_status: Some(status),
+            }
+        }
+    }
+
+    impl<B> Service<Request<B>> for CaptureService {
+        type Response = http::Response<EmptyBody>;
+        type Error = Infallible;
+        type Future = std::future::Ready<Result<Self::Response, Self::Error>>;
+
+        fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+            Poll::Ready(Ok(()))
+        }
+
+        fn call(&mut self, req: Request<B>) -> Self::Future {
+            *self.seen.lock().unwrap() = Some(req.headers().clone());
+            let mut builder = http::Response::builder().status(200);
+            if let Some(status) = self.response_grpc_status {
+                builder = builder.header("grpc-status", status);
+            }
+            std::future::ready(Ok(builder.body(Empty::new()).unwrap()))
+        }
+    }
+
+    fn post_req() -> Request<EmptyBody> {
+        Request::builder()
+            .method("POST")
+            .uri("http://svc/package.Service/Method")
+            .body(Empty::new())
+            .unwrap()
+    }
+
+    // With no dispatcher active, the client span must not be created —
+    // `tracing`'s `log` compat would downgrade it into `grpc_request; ...`
+    // log spam in processes that only configure a `log` logger.
+    // See `crate::dispatcher_active`.
+    #[test]
+    fn make_client_span_without_dispatcher_is_none() {
+        assert!(MakeClientSpan.make_span(&post_req()).is_none());
+    }
+
+    #[test]
+    fn make_client_span_with_scoped_dispatcher_is_enabled() {
+        let _env = OtelTestEnv::install();
+        assert!(!MakeClientSpan.make_span(&post_req()).is_disabled());
+    }
+
+    #[tokio::test]
+    async fn inject_under_trace_layer_uses_client_span_not_parent() {
+        let _env = OtelTestEnv::install();
+
+        let capture = CaptureService::new();
+        let seen = Arc::clone(&capture.seen);
+        let mut svc = ServiceBuilder::new()
+            .layer(TraceLayer::new_for_grpc().make_span_with(MakeClientSpan))
+            .layer(InjectTraceContextLayer)
+            .service(capture);
+
+        let parent = tracing::info_span!("parent_handler");
+        let parent_span_id = otel_span_id_hex(&parent);
+        let parent_trace_id = otel_trace_id_hex(&parent);
+        assert_ne!(parent_span_id, "0000000000000000");
+
+        async {
+            let fut = Service::call(&mut svc, post_req());
+            fut.await.unwrap();
+        }
+        .instrument(parent)
+        .await;
+
+        let headers = seen.lock().unwrap().clone().expect("headers");
+        let tp = headers
+            .get("traceparent")
+            .expect("traceparent")
+            .to_str()
+            .unwrap();
+        let (_ver, injected_trace_id, injected_span_id) = parse_traceparent(tp);
+
+        assert_eq!(injected_trace_id, parent_trace_id);
+        assert_ne!(injected_span_id, parent_span_id);
+    }
+
+    #[tokio::test]
+    async fn inject_without_trace_layer_uses_parent_span() {
+        let _env = OtelTestEnv::install();
+
+        let capture = CaptureService::new();
+        let seen = Arc::clone(&capture.seen);
+        let mut svc = ServiceBuilder::new()
+            .layer(InjectTraceContextLayer)
+            .service(capture);
+
+        let parent = tracing::info_span!("parent_handler");
+        let parent_span_id = otel_span_id_hex(&parent);
+        let parent_trace_id = otel_trace_id_hex(&parent);
+
+        async {
+            let fut = Service::call(&mut svc, post_req());
+            fut.await.unwrap();
+        }
+        .instrument(parent)
+        .await;
+
+        let headers = seen.lock().unwrap().clone().expect("headers captured");
+        let tp = headers.get("traceparent").unwrap().to_str().unwrap();
+        let (_ver, injected_trace_id, injected_span_id) = parse_traceparent(tp);
+
+        assert_eq!(injected_trace_id, parent_trace_id);
+        assert_eq!(injected_span_id, parent_span_id);
+    }
+
+    #[tokio::test]
+    async fn grpc_status_non_ok_invokes_on_failure_classifier() {
+        let _env = OtelTestEnv::install();
+
+        let failures: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
+        let failures_cb = Arc::clone(&failures);
+
+        let capture = CaptureService::with_grpc_status("13");
+        let mut svc = ServiceBuilder::new()
+            .layer(
+                TraceLayer::new_for_grpc()
+                    .make_span_with(MakeClientSpan)
+                    .on_failure(
+                        move |class: GrpcFailureClass,
+                              _latency: Duration,
+                              _span: &tracing::Span| {
+                            failures_cb.lock().unwrap().push(class.to_string());
+                        },
+                    ),
+            )
+            .layer(InjectTraceContextLayer)
+            .service(capture);
+
+        let fut = Service::call(&mut svc, post_req());
+        let _ = fut.await.unwrap();
+
+        let recorded = failures.lock().unwrap().clone();
+        assert_eq!(recorded.len(), 1, "{recorded:?}");
+        assert!(
+            recorded[0].contains("13") || recorded[0].to_lowercase().contains("code"),
+            "{}",
+            recorded[0]
+        );
+    }
+
+    #[tokio::test]
+    async fn grpc_status_ok_does_not_invoke_on_failure() {
+        let _env = OtelTestEnv::install();
+
+        let failures: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
+        let failures_cb = Arc::clone(&failures);
+
+        let capture = CaptureService::with_grpc_status("0");
+        let mut svc = ServiceBuilder::new()
+            .layer(
+                TraceLayer::new_for_grpc()
+                    .make_span_with(MakeClientSpan)
+                    .on_failure(
+                        move |class: GrpcFailureClass,
+                              _latency: Duration,
+                              _span: &tracing::Span| {
+                            failures_cb.lock().unwrap().push(class.to_string());
+                        },
+                    ),
+            )
+            .layer(InjectTraceContextLayer)
+            .service(capture);
+
+        let fut = Service::call(&mut svc, post_req());
+        let _ = fut.await.unwrap();
+
+        assert!(failures.lock().unwrap().is_empty());
+    }
+
+    #[tokio::test]
+    async fn attach_trace_to_grpc_request_sets_traceparent_metadata() {
+        let _env = OtelTestEnv::install();
+        let span = tracing::info_span!("handler");
+        let span_id = otel_span_id_hex(&span);
+        let _enter = span.enter();
+
+        let req = attach_trace_to_grpc_request(tonic::Request::new(())).unwrap();
+        let tp = req
+            .metadata()
+            .get("traceparent")
+            .expect("traceparent in metadata")
+            .to_str()
+            .unwrap();
+        let (_ver, _tid, injected_span_id) = parse_traceparent(tp);
+        assert_eq!(injected_span_id, span_id);
+    }
+
+    #[tokio::test]
+    async fn make_client_span_records_otel_kind_client() {
+        let env = OtelTestEnv::install();
+        {
+            let req = post_req();
+            let mut make = MakeClientSpan;
+            let span = make.make_span(&req);
+            let _e = span.enter();
+        }
+        let spans = env.finished_spans();
+        let grpc = spans
+            .iter()
+            .find(|s| s.name == "grpc_request")
+            .expect("grpc_request");
+        assert_eq!(grpc.span_kind, opentelemetry::trace::SpanKind::Client);
+    }
+}

crates/common/xai-tracing/src/http_client.rs

@@ -0,0 +1,155 @@
+use async_trait::async_trait;
+use opentelemetry::global;
+use opentelemetry_http::HeaderInjector;
+use reqwest::header::HeaderMap;
+use reqwest_middleware::{ClientBuilder, ClientWithMiddleware, Middleware};
+use tracing::{Instrument, Span, field};
+use tracing_opentelemetry::OpenTelemetrySpanExt;
+
+pub fn attach_trace_to_http_request(headers: &mut HeaderMap) {
+    global::get_text_map_propagator(|propagator| {
+        let context = Span::current().context();
+        propagator.inject_context(&context, &mut HeaderInjector(headers));
+    });
+}
+
+pub type TracedHttpClient = ClientWithMiddleware;
+
+pub fn traced_client(client: reqwest::Client) -> TracedHttpClient {
+    ClientBuilder::new(client).with(TracingMiddleware).build()
+}
+
+pub fn traced_client_new() -> TracedHttpClient {
+    traced_client(reqwest::Client::new())
+}
+
+pub fn traced_client_from_builder(
+    builder: reqwest::ClientBuilder,
+) -> Result<TracedHttpClient, reqwest::Error> {
+    Ok(traced_client(builder.build()?))
+}
+
+#[derive(Clone, Debug, Default)]
+struct TracingMiddleware;
+
+#[async_trait]
+impl Middleware for TracingMiddleware {
+    async fn handle(
+        &self,
+        mut req: reqwest::Request,
+        extensions: &mut http::Extensions,
+        next: reqwest_middleware::Next<'_>,
+    ) -> reqwest_middleware::Result<reqwest::Response> {
+        let method = req.method().as_str().to_owned();
+        let url = req.url().clone();
+        // No active dispatcher → the span has no consumer and would only be
+        // downgraded to `log` spam. See `crate::dispatcher_active`.
+        let span = if crate::dispatcher_active() {
+            tracing::info_span!(
+                "http_request",
+                otel.kind = "client",
+                "http.request.method" = %method,
+                "url.full" = %url,
+                "http.response.status_code" = field::Empty,
+            )
+        } else {
+            Span::none()
+        };
+
+        let result = async move {
+            attach_trace_to_http_request(req.headers_mut());
+            next.run(req, extensions).await
+        }
+        .instrument(span.clone())
+        .await;
+
+        if let Ok(ref response) = result {
+            span.record("http.response.status_code", response.status().as_u16());
+        }
+        result
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::testing::{OtelTestEnv, otel_span_id_hex, otel_trace_id_hex, parse_traceparent};
+    use tracing::Instrument;
+    use wiremock::matchers::method;
+    use wiremock::{Mock, MockServer, ResponseTemplate};
+
+    #[tokio::test]
+    async fn attach_trace_to_http_request_writes_traceparent() {
+        let _env = OtelTestEnv::install();
+        let span = tracing::info_span!("http_request", otel.kind = "client");
+        let span_id = otel_span_id_hex(&span);
+        let mut headers = HeaderMap::new();
+        let _enter = span.enter();
+        attach_trace_to_http_request(&mut headers);
+        let tp = headers.get("traceparent").unwrap().to_str().unwrap();
+        let (_ver, _tid, injected_span_id) = parse_traceparent(tp);
+        assert_eq!(injected_span_id, span_id);
+    }
+
+    #[tokio::test]
+    async fn traced_client_injects_client_span_not_parent_on_wire() {
+        let _env = OtelTestEnv::install();
+        let server = MockServer::start().await;
+        Mock::given(method("GET"))
+            .respond_with(ResponseTemplate::new(200))
+            .mount(&server)
+            .await;
+
+        let client = traced_client(reqwest::Client::new());
+        let parent = tracing::info_span!("parent_handler");
+        let parent_span_id = otel_span_id_hex(&parent);
+        let parent_trace_id = otel_trace_id_hex(&parent);
+
+        async {
+            client
+                .get(format!("{}/health", server.uri()))
+                .send()
+                .await
+                .unwrap();
+        }
+        .instrument(parent)
+        .await;
+
+        let received = server.received_requests().await.unwrap();
+        assert_eq!(received.len(), 1);
+        let tp = received[0]
+            .headers
+            .get("traceparent")
+            .expect("traceparent on wire")
+            .to_str()
+            .unwrap();
+        let (_ver, injected_trace_id, injected_span_id) = parse_traceparent(tp);
+
+        assert_eq!(injected_trace_id, parent_trace_id);
+        assert_ne!(injected_span_id, parent_span_id);
+    }
+
+    #[tokio::test]
+    async fn traced_client_returns_response_status() {
+        let _env = OtelTestEnv::install();
+        let server = MockServer::start().await;
+        Mock::given(method("GET"))
+            .respond_with(ResponseTemplate::new(404))
+            .mount(&server)
+            .await;
+
+        let client = traced_client_new();
+        let resp = client
+            .get(format!("{}/missing", server.uri()))
+            .send()
+            .await
+            .unwrap();
+        assert_eq!(resp.status().as_u16(), 404);
+        assert!(
+            server.received_requests().await.unwrap()[0]
+                .headers
+                .get("traceparent")
+                .is_some()
+        );
+    }
+}

crates/common/xai-tracing/src/lib.rs

@@ -0,0 +1,19 @@
+mod dispatch;
+mod grpc_client;
+mod timer;
+
+pub mod fastrace;
+pub mod http_client;
+pub mod tokio;
+
+#[cfg(test)]
+mod testing;
+
+pub use dispatch::*;
+pub use fastrace::*;
+pub use grpc_client::*;
+pub use http_client::{
+    TracedHttpClient, attach_trace_to_http_request, traced_client, traced_client_from_builder,
+    traced_client_new,
+};
+pub use timer::*;

crates/common/xai-tracing/src/testing.rs

@@ -0,0 +1,65 @@
+use opentelemetry::global;
+use opentelemetry::trace::{SpanContext, TraceContextExt, TracerProvider as _};
+use opentelemetry_sdk::propagation::TraceContextPropagator;
+use opentelemetry_sdk::trace::{
+    InMemorySpanExporter, InMemorySpanExporterBuilder, SdkTracerProvider, SimpleSpanProcessor,
+};
+use tracing::Span;
+use tracing_opentelemetry::OpenTelemetrySpanExt;
+use tracing_subscriber::prelude::*;
+
+pub struct OtelTestEnv {
+    _guard: tracing::subscriber::DefaultGuard,
+    provider: SdkTracerProvider,
+    exporter: InMemorySpanExporter,
+}
+
+impl OtelTestEnv {
+    pub fn install() -> Self {
+        global::set_text_map_propagator(TraceContextPropagator::new());
+        let exporter = InMemorySpanExporterBuilder::new().build();
+        let provider = SdkTracerProvider::builder()
+            .with_span_processor(SimpleSpanProcessor::new(exporter.clone()))
+            .build();
+        let tracer = provider.tracer("xai-tracing-test");
+        let otel_layer = tracing_opentelemetry::layer()
+            .with_tracer(tracer)
+            .with_context_activation(false)
+            .with_filter(tracing_subscriber::filter::LevelFilter::INFO);
+        let guard = tracing_subscriber::registry()
+            .with(otel_layer)
+            .set_default();
+        Self {
+            _guard: guard,
+            provider,
+            exporter,
+        }
+    }
+
+    pub fn finished_spans(&self) -> Vec<opentelemetry_sdk::trace::SpanData> {
+        let _ = self.provider.force_flush();
+        self.exporter.get_finished_spans().unwrap_or_default()
+    }
+}
+
+pub fn parse_traceparent(value: &str) -> (&str, &str, &str) {
+    let mut parts = value.split('-');
+    let version = parts.next().expect("version");
+    let trace_id = parts.next().expect("trace_id");
+    let span_id = parts.next().expect("span_id");
+    (version, trace_id, span_id)
+}
+
+pub fn otel_span_id_hex(span: &Span) -> String {
+    let cx = span.context();
+    let span_ref = cx.span();
+    let sc: &SpanContext = span_ref.span_context();
+    format!("{:016x}", u64::from_be_bytes(sc.span_id().to_bytes()))
+}
+
+pub fn otel_trace_id_hex(span: &Span) -> String {
+    let cx = span.context();
+    let span_ref = cx.span();
+    let sc: &SpanContext = span_ref.span_context();
+    format!("{:032x}", u128::from_be_bytes(sc.trace_id().to_bytes()))
+}

crates/common/xai-tracing/src/timer.rs

@@ -0,0 +1,62 @@
+use log::error;
+use log::info;
+use tokio::time::Instant;
+
+/// A simple timer that logs the runtime of an operation.
+pub struct Timer {
+    /// Time the operation started.
+    start: Instant,
+    /// An ID shown in the logs to associate the log messages from the timer with each other-
+    id: uuid::Uuid,
+    /// A string that is being logged.
+    message: String,
+    /// True if the timer has been stopped already.
+    stopped: bool,
+}
+
+impl Timer {
+    /// Creates a new Timer instance and starts the timer.
+    pub fn new<S: AsRef<str>>(message: S) -> Self {
+        let id = uuid::Uuid::new_v4();
+        info!("[{}] START: {}", id, message.as_ref());
+        Self {
+            start: Instant::now(),
+            id,
+            message: message.as_ref().to_string(),
+            stopped: false,
+        }
+    }
+
+    /// Stops the timer and logs the result.
+    pub fn stop<T>(&mut self, result: T) -> T {
+        if !self.stopped {
+            let runtime = self.start.elapsed().as_secs_f32();
+            info!(
+                "[{}] FINISHED in {:.3}s: {}",
+                self.id, runtime, self.message
+            );
+            self.stopped = true;
+        }
+
+        result
+    }
+
+    /// Stop the timer prematurely and logs an error.
+    pub fn force_stop(&mut self) {
+        if !self.stopped {
+            let runtime = self.start.elapsed().as_secs_f32();
+            error!(
+                "[{}] FAILED after {:.3}s: {}",
+                self.id, runtime, self.message
+            );
+            self.stopped = true;
+        }
+    }
+}
+
+/// Automatically report the runtime when the object is dropped.
+impl Drop for Timer {
+    fn drop(&mut self) {
+        self.force_stop();
+    }
+}

crates/common/xai-tracing/src/tokio.rs

@@ -0,0 +1,22 @@
+use std::future::Future;
+use tokio::task::JoinHandle;
+use tracing::{Instrument, Span};
+
+/// Utility macro for propagating the current tracing context to a newly spawned task.
+///
+/// Note: The spawned task will be associated with the currently active span. To create a *new*
+/// span for the spawned task, manually instrument the future using [tracing::Instrument] instead of
+/// using this macro. For example:
+///
+/// use tracing::{info_span, Instrument};
+///
+/// let fut = tokio::spawn(async move {
+///     print!("do stuff")
+/// }.instrument(info_span!("spawned task")));
+pub fn spawn_traced<F>(future: F) -> JoinHandle<F::Output>
+where
+    F: Future + Send + 'static,
+    F::Output: Send + 'static,
+{
+    tokio::spawn(future.instrument(Span::current()))
+}