stasis-rs 0.9.3

Durable AI orchestration framework with runtime jobs, lineage, and memory integration
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

use axum::extract::Request;
use axum::http::HeaderMap;
use axum::middleware::Next;
use axum::response::Response;

use crate::application::telemetry::propagation::parse_traceparent;
use crate::application::telemetry::request_context;
use crate::ports::outbound::runtime::runtime_tracing::TraceContext;

pub fn extract_traceparent_from_headers(headers: &HeaderMap) -> Option<TraceContext> {
    headers
        .get("traceparent")
        .or_else(|| headers.get("Traceparent"))
        .and_then(|value| value.to_str().ok())
        .and_then(|header| parse_traceparent(header).ok())
}

pub async fn propagate_inbound_trace_context(request: Request, next: Next) -> Response {
    let Some(trace) = extract_traceparent_from_headers(request.headers()) else {
        return next.run(request).await;
    };

    request_context::scope_inbound_trace(trace, async move { next.run(request).await }).await
}

#[cfg(test)]
mod tests {
    use axum::http::HeaderMap;

    use super::*;

    #[test]
    fn extract_traceparent_from_headers_is_case_insensitive_on_key() {
        let mut headers = HeaderMap::new();
        headers.insert(
            "traceparent",
            "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
                .parse()
                .expect("valid header value"),
        );

        let context = extract_traceparent_from_headers(&headers).expect("trace context");
        assert_eq!(context.trace_id, "4bf92f3577b34da6a3ce929d0e0e4736");
    }
}