Skip to main content

systemprompt_api/services/middleware/
trace.rs

1//! Trace-id propagation middleware.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use axum::extract::Request;
7use axum::middleware::Next;
8use axum::response::Response;
9use systemprompt_models::RequestContext;
10
11pub async fn inject_trace_header(request: Request, next: Next) -> Response {
12    let trace_id = request
13        .extensions()
14        .get::<RequestContext>()
15        .map(|ctx| ctx.trace_id().as_str().to_owned());
16
17    let mut response = next.run(request).await;
18
19    if let Some(id) = trace_id
20        && let Ok(header_value) = id.parse()
21    {
22        response.headers_mut().insert("x-trace-id", header_value);
23    }
24
25    response
26}