34 lines
1.1 KiB
Rust
34 lines
1.1 KiB
Rust
use axum::{extract::Request, http::HeaderValue, middleware::Next, response::Response};
|
|
use crank_observability::{RequestId, set_remote_trace_parent, with_request_correlation};
|
|
use tracing::{Instrument, info_span};
|
|
|
|
use crate::transport::HEADER_X_REQUEST_ID;
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub(super) struct RequestContext {
|
|
pub(super) request_id: String,
|
|
}
|
|
|
|
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 span = info_span!(
|
|
target: "crank::trace",
|
|
"mcp.request",
|
|
request_id = %request_id,
|
|
);
|
|
set_remote_trace_parent(&span, request.headers());
|
|
request.extensions_mut().insert(context);
|
|
|
|
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
|
|
})
|
|
.await
|
|
}
|