fix(lambda-http): source X-Ray trace id from context

alessandrobologna · alessandrobologna · commit 949c00231503 · 2025-12-23T15:39:06.000-05:00
- Replace env-var based X-Ray header injection with Context.xray_trace_id
- Set x-amzn-trace-id in HTTP and streaming adapters before handler call
- Remove legacy update_xray_trace_id_header helper from request parsing

Tests:
- cargo +nightly fmt --all
- cargo clippy --workspace --all-features -- -D warnings
- cargo +stable test --all-features -p lambda_runtime_api_client -p lambda_runtime -p lambda_http
diff --git a/lambda-http/src/lib.rs b/lambda-http/src/lib.rs
@@ -192,9 +192,11 @@ where
     }
 
     fn call(&mut self, req: LambdaEvent<LambdaRequest>) -> Self::Future {
-        let request_origin = req.payload.request_origin();
-        let event: Request = req.payload.into();
-        let fut = Box::pin(self.service.call(event.with_lambda_context(req.context)));
+        let LambdaEvent { payload, context } = req;
+        let request_origin = payload.request_origin();
+        let mut event: Request = payload.into();
+        update_xray_trace_id_header_from_context(event.headers_mut(), &context);
+        let fut = Box::pin(self.service.call(event.with_lambda_context(context)));
 
         TransformResponse::Request(request_origin, fut)
     }
@@ -233,6 +235,15 @@ where
     lambda_runtime::run_concurrent(Adapter::from(handler)).await
 }
 
+// Replaces update_xray_trace_id_header (env var), now set from Context
+fn update_xray_trace_id_header_from_context(headers: &mut http::HeaderMap, context: &Context) {
+    if let Some(trace_id) = context.xray_trace_id.as_deref() {
+        if let Ok(header_value) = http::HeaderValue::from_str(trace_id) {
+            headers.insert(http::header::HeaderName::from_static("x-amzn-trace-id"), header_value);
+        }
+    }
+}
+
 #[cfg(test)]
 mod test_adapter {
     use std::task::{Context, Poll};
diff --git a/lambda-http/src/request.rs b/lambda-http/src/request.rs
@@ -146,7 +146,6 @@ fn into_api_gateway_v2_request(ag: ApiGatewayV2httpRequest) -> http::Request<Bod
         .extension(RequestContext::ApiGatewayV2(ag.request_context));
 
     let mut headers = ag.headers;
-    update_xray_trace_id_header(&mut headers);
     if let Some(cookies) = ag.cookies {
         if let Ok(header_value) = HeaderValue::from_str(&cookies.join(";")) {
             headers.insert(http::header::COOKIE, header_value);
@@ -170,13 +169,6 @@ fn into_api_gateway_v2_request(ag: ApiGatewayV2httpRequest) -> http::Request<Bod
     req
 }
 
-fn update_xray_trace_id_header(headers: &mut HeaderMap) {
-    if let Ok(xray_trace_id) = env::var("_X_AMZN_TRACE_ID") {
-        if let Ok(header_value) = HeaderValue::from_str(&xray_trace_id) {
-            headers.insert(HeaderName::from_static("x-amzn-trace-id"), header_value);
-        }
-    }
-}
 #[cfg(feature = "apigw_rest")]
 fn into_proxy_request(ag: ApiGatewayProxyRequest) -> http::Request<Body> {
     let http_method = ag.http_method;
@@ -214,7 +206,6 @@ fn into_proxy_request(ag: ApiGatewayProxyRequest) -> http::Request<Body> {
     // multi-value_headers our cannoncial source of request headers
     let mut headers = ag.multi_value_headers;
     headers.extend(ag.headers);
-    update_xray_trace_id_header(&mut headers);
 
     let base64 = ag.is_base64_encoded;
     let mut req = builder
@@ -265,7 +256,6 @@ fn into_alb_request(alb: AlbTargetGroupRequest) -> http::Request<Body> {
     // multi-value_headers our cannoncial source of request headers
     let mut headers = alb.multi_value_headers;
     headers.extend(alb.headers);
-    update_xray_trace_id_header(&mut headers);
 
     let base64 = alb.is_base64_encoded;
 
@@ -330,7 +320,6 @@ fn into_websocket_request(ag: ApiGatewayWebsocketProxyRequest) -> http::Request<
     // multi-value_headers our canonical source of request headers
     let mut headers = ag.multi_value_headers;
     headers.extend(ag.headers);
-    update_xray_trace_id_header(&mut headers);
 
     let base64 = ag.is_base64_encoded;
     let mut req = builder
@@ -355,8 +344,6 @@ fn into_pass_through_request(data: String) -> http::Request<Body> {
     let headers = builder.headers_mut().unwrap();
     headers.insert("Content-Type", "application/json".parse().unwrap());
 
-    update_xray_trace_id_header(headers);
-
     let raw_path = "/events";
 
     builder
diff --git a/lambda-http/src/streaming.rs b/lambda-http/src/streaming.rs
@@ -1,4 +1,6 @@
-use crate::{http::header::SET_COOKIE, request::LambdaRequest, Request, RequestExt};
+use crate::{
+    http::header::SET_COOKIE, request::LambdaRequest, update_xray_trace_id_header_from_context, Request, RequestExt,
+};
 use bytes::Bytes;
 use core::{
     fmt::Debug,
@@ -60,10 +62,12 @@ where
     }
 
     fn call(&mut self, req: LambdaEvent<LambdaRequest>) -> Self::Future {
-        let event: Request = req.payload.into();
+        let LambdaEvent { payload, context } = req;
+        let mut event: Request = payload.into();
+        update_xray_trace_id_header_from_context(event.headers_mut(), &context);
         Box::pin(
             self.service
-                .call(event.with_lambda_context(req.context))
+                .call(event.with_lambda_context(context))
                 .map_ok(into_stream_response),
         )
     }
@@ -148,8 +152,10 @@ where
 }
 
 fn event_to_request(req: LambdaEvent<LambdaRequest>) -> Request {
-    let event: Request = req.payload.into();
-    event.with_lambda_context(req.context)
+    let LambdaEvent { payload, context } = req;
+    let mut event: Request = payload.into();
+    update_xray_trace_id_header_from_context(event.headers_mut(), &context);
+    event.with_lambda_context(context)
 }
 
 /// Runs the Lambda runtime with a handler that returns **streaming** HTTP
