crank/apps/mcp-server/tests/integration/request_context.rs

use std::{
    sync::{Arc, Mutex},
    time::Duration,
};

use axum::{
    body::Body,
    http::{Request, StatusCode},
};
use opentelemetry::{
    global,
    trace::{TraceId, TracerProvider as _},
};
use opentelemetry_sdk::{
    error::OTelSdkResult,
    propagation::TraceContextPropagator,
    trace::{SdkTracerProvider, SpanData, SpanExporter},
};
use tower::ServiceExt;
use tracing::instrument::WithSubscriber;
use tracing_subscriber::layer::SubscriberExt;

use super::common::{build_test_app, test_registry};

const REMOTE_TRACE_ID: &str = "0af7651916cd43dd8448eb211c80319c";

#[tokio::test(flavor = "current_thread")]
async fn covers_valid_invalid_and_absent_traceparent_on_mcp_boundary() {
    global::set_text_map_propagator(TraceContextPropagator::new());
    let exported = Arc::new(Mutex::new(Vec::new()));
    let provider = SdkTracerProvider::builder()
        .with_simple_exporter(CapturingExporter(Arc::clone(&exported)))
        .build();
    let tracer = provider.tracer("mcp-request-context-test");
    let subscriber =
        tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer));
    let dispatch = tracing::Dispatch::new(subscriber);
    let app = build_test_app(test_registry().await, Duration::ZERO, None);

    let (valid, invalid, absent) = async {
        let valid = send_health(
            app.clone(),
            Some("00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"),
            Some("request-id-is-separate"),
        )
        .await;
        let invalid = send_health(
            app.clone(),
            Some("canary-invalid-traceparent"),
            Some("bad,value"),
        )
        .await;
        let absent = send_health(app, None, None).await;
        (valid, invalid, absent)
    }
    .with_subscriber(dispatch)
    .await;
    provider.force_flush().unwrap();

    assert_eq!(valid.status, StatusCode::OK);
    assert_eq!(valid.request_id.as_deref(), Some("request-id-is-separate"));
    assert_eq!(invalid.status, StatusCode::OK);
    assert_eq!(absent.status, StatusCode::OK);
    assert!(valid.traceparent_response.is_none());
    assert!(invalid.traceparent_response.is_none());
    assert!(absent.traceparent_response.is_none());

    let trace_ids: Vec<_> = exported
        .lock()
        .unwrap()
        .iter()
        .filter(|span| span.name.as_ref() == "mcp.request")
        .map(|span| span.span_context.trace_id())
        .collect();
    assert_eq!(trace_ids.len(), 3);
    assert_eq!(trace_ids[0].to_string(), REMOTE_TRACE_ID);
    assert_ne!(trace_ids[1], trace_ids[0]);
    assert_ne!(trace_ids[2], trace_ids[0]);
    assert_ne!(trace_ids[1], trace_ids[2]);
    assert!(!trace_ids.contains(&TraceId::INVALID));
    provider.shutdown().unwrap();
}

async fn send_health(
    app: axum::Router,
    traceparent: Option<&str>,
    request_id: Option<&str>,
) -> ProbeResponse {
    let mut request = Request::builder().uri("/health");
    if let Some(traceparent) = traceparent {
        request = request.header("traceparent", traceparent);
    }
    if let Some(request_id) = request_id {
        request = request.header("x-request-id", request_id);
    }
    let response = app
        .oneshot(request.body(Body::empty()).unwrap())
        .await
        .unwrap();

    ProbeResponse {
        status: response.status(),
        request_id: response
            .headers()
            .get("x-request-id")
            .and_then(|value| value.to_str().ok())
            .map(str::to_owned),
        traceparent_response: response
            .headers()
            .get("traceparent")
            .and_then(|value| value.to_str().ok())
            .map(str::to_owned),
    }
}

struct ProbeResponse {
    status: StatusCode,
    request_id: Option<String>,
    traceparent_response: Option<String>,
}

#[derive(Clone, Debug)]
struct CapturingExporter(Arc<Mutex<Vec<SpanData>>>);

impl SpanExporter for CapturingExporter {
    async fn export(&self, batch: Vec<SpanData>) -> OTelSdkResult {
        self.0.lock().unwrap().extend(batch);
        Ok(())
    }
}