feat: harden community production foundation through story 1.5

crates/crank-community-mcp/src/request_context.rs

@@ -1,33 +1,101 @@
 use axum::{extract::Request, http::HeaderValue, middleware::Next, response::Response};
-use crank_observability::{RequestId, set_remote_trace_parent, with_request_correlation};
+use crank_core::{CorrelationContext, RequestId, TraceContext};
+use crank_observability::{set_remote_trace_parent, with_request_correlation};
 use tracing::{Instrument, info_span};
 
 use crate::transport::HEADER_X_REQUEST_ID;
 
+const HEADER_X_TRACE_ID: axum::http::HeaderName = axum::http::HeaderName::from_static("x-trace-id");
+
 #[derive(Clone, Debug)]
 pub(super) struct RequestContext {
-    pub(super) request_id: String,
+    pub(super) correlation: CorrelationContext,
+}
+
+impl RequestContext {
+    pub(super) fn request_id(&self) -> &str {
+        self.correlation.request_id().as_str()
+    }
 }
 
 pub(super) async fn apply_request_context(mut request: Request, next: Next) -> Response {
-    let request_id = RequestId::resolve_from_headers(request.headers()).into_string();
-    let context = RequestContext {
-        request_id: request_id.clone(),
-    };
+    let (request_id, remote_parent) = resolve_correlation(request.headers());
     let span = info_span!(
         target: "crank::trace",
         "mcp.request",
         request_id = %request_id,
+        trace_id = tracing::field::Empty,
     );
-    set_remote_trace_parent(&span, request.headers());
-    request.extensions_mut().insert(context);
+    if let Some(remote_parent) = remote_parent.as_ref() {
+        set_canonical_parent(&span, remote_parent);
+    }
+    let trace_context = crank_trace::trace_context_for_span(&span).unwrap_or_else(|| {
+        remote_parent
+            .as_ref()
+            .map_or_else(TraceContext::generate, TraceContext::continue_local)
+    });
+    span.record("trace_id", trace_context.trace_id().as_str());
+    let context = RequestContext {
+        correlation: CorrelationContext::new(request_id, trace_context),
+    };
+    request.extensions_mut().insert(context.clone());
 
-    with_request_correlation(request_id.clone(), async move {
-        let mut response = next.run(request).instrument(span).await;
-        if let Ok(value) = HeaderValue::from_str(&request_id) {
-            response.headers_mut().insert(HEADER_X_REQUEST_ID, value);
-        }
-        response
-    })
+    with_request_correlation(
+        context.correlation.request_id().to_string(),
+        context.correlation.trace_id().to_string(),
+        async move {
+            let mut response = next.run(request).instrument(span).await;
+            if let Ok(value) = HeaderValue::from_str(context.correlation.request_id().as_str()) {
+                response.headers_mut().insert(HEADER_X_REQUEST_ID, value);
+            }
+            if let Ok(value) = HeaderValue::from_str(context.correlation.trace_id().as_str()) {
+                response.headers_mut().insert(HEADER_X_TRACE_ID, value);
+            }
+            response
+        },
+    )
     .await
 }
+
+fn resolve_correlation(headers: &axum::http::HeaderMap) -> (RequestId, Option<TraceContext>) {
+    let _tracestate_accepted = one_auxiliary_header_within_budget(
+        headers,
+        "tracestate",
+        TraceContext::tracestate_within_budget,
+    );
+    let _baggage_accepted =
+        one_auxiliary_header_within_budget(headers, "baggage", TraceContext::baggage_within_budget);
+    let mut request_ids = headers.get_all(HEADER_X_REQUEST_ID).iter();
+    let request_id = request_ids.next().and_then(|value| value.to_str().ok());
+    let request_id = if request_ids.next().is_some() {
+        RequestId::generate()
+    } else {
+        RequestId::resolve(request_id)
+    };
+    let mut traceparents = headers.get_all("traceparent").iter();
+    let traceparent = traceparents.next().and_then(|value| value.to_str().ok());
+    let remote_parent = if traceparents.next().is_some() {
+        None
+    } else {
+        traceparent.and_then(|value| TraceContext::parse(value).ok())
+    };
+    (request_id, remote_parent)
+}
+
+fn one_auxiliary_header_within_budget(
+    headers: &axum::http::HeaderMap,
+    name: &'static str,
+    validate: fn(&str) -> bool,
+) -> bool {
+    let mut values = headers.get_all(name).iter();
+    let value = values.next().and_then(|value| value.to_str().ok());
+    values.next().is_none() && value.is_some_and(validate)
+}
+
+fn set_canonical_parent(span: &tracing::Span, context: &TraceContext) {
+    let mut headers = axum::http::HeaderMap::new();
+    if let Ok(value) = HeaderValue::from_str(context.traceparent()) {
+        headers.insert("traceparent", value);
+        set_remote_trace_parent(span, &headers);
+    }
+}