LLMO-1023: add trace ID propagation for SQS messages and HTTP headers

packages/spacecat-shared-http-utils/src/enrich-path-info-wrapper.js

@@ -13,14 +13,23 @@
 export function enrichPathInfo(fn) { // export for testing
   return async (request, context) => {
     const [, route] = context?.pathInfo?.suffix?.split(/\/+/) || [];
+    const headers = request.headers.plain();
+
     context.pathInfo = {
       ...context.pathInfo,
       ...{
         method: request.method.toUpperCase(),
-        headers: request.headers.plain(),
+        headers,
         route,
       },
     };
+
+    // Extract and store traceId from x-trace-id header if present
+    const traceIdHeader = headers['x-trace-id'];
+    if (traceIdHeader) {
+      context.traceId = traceIdHeader;
+    }
+
     return fn(request, context);
   };
 }

packages/spacecat-shared-http-utils/test/enrich-path-info-wrapper.test.js

@@ -0,0 +1,150 @@
+/*
+ * Copyright 2025 Adobe. All rights reserved.
+ * This file is licensed to you under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License. You may obtain a copy
+ * of the License at http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under
+ * the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS
+ * OF ANY KIND, either express or implied. See the License for the specific language
+ * governing permissions and limitations under the License.
+ */
+/* eslint-env mocha */
+import { expect } from 'chai';
+import sinon from 'sinon';
+
+import { enrichPathInfo } from '../src/enrich-path-info-wrapper.js';
+
+describe('enrichPathInfo', () => {
+  let mockRequest;
+  let mockContext;
+  let mockFn;
+
+  beforeEach(() => {
+    mockFn = sinon.stub().resolves({ status: 200 });
+
+    mockRequest = {
+      method: 'POST',
+      headers: {
+        plain: () => ({
+          'content-type': 'application/json',
+          'user-agent': 'test-agent',
+        }),
+      },
+    };
+
+    mockContext = {
+      pathInfo: {
+        suffix: '/api/test',
+      },
+    };
+  });
+
+  afterEach(() => {
+    sinon.restore();
+  });
+
+  it('should enrich context with pathInfo including method, headers, and route', async () => {
+    const wrapper = enrichPathInfo(mockFn);
+    await wrapper(mockRequest, mockContext);
+
+    expect(mockContext.pathInfo).to.deep.include({
+      method: 'POST',
+      route: 'api',
+    });
+    expect(mockContext.pathInfo.headers).to.deep.equal({
+      'content-type': 'application/json',
+      'user-agent': 'test-agent',
+    });
+  });
+
+  it('should extract traceId from x-trace-id header and store in context', async () => {
+    mockRequest.headers.plain = () => ({
+      'content-type': 'application/json',
+      'x-trace-id': '1-5e8e8e8e-5e8e8e8e5e8e8e8e5e8e8e8e',
+    });
+
+    const wrapper = enrichPathInfo(mockFn);
+    await wrapper(mockRequest, mockContext);
+
+    expect(mockContext.traceId).to.equal('1-5e8e8e8e-5e8e8e8e5e8e8e8e5e8e8e8e');
+  });
+
+  it('should not set traceId in context when x-trace-id header is missing', async () => {
+    mockRequest.headers.plain = () => ({
+      'content-type': 'application/json',
+    });
+
+    const wrapper = enrichPathInfo(mockFn);
+    await wrapper(mockRequest, mockContext);
+
+    expect(mockContext.traceId).to.be.undefined;
+  });
+
+  it('should handle case-sensitive x-trace-id header', async () => {
+    mockRequest.headers.plain = () => ({
+      'content-type': 'application/json',
+      'X-Trace-Id': '1-different-case',
+    });
+
+    const wrapper = enrichPathInfo(mockFn);
+    await wrapper(mockRequest, mockContext);
+
+    // Header keys should be lowercase
+    expect(mockContext.traceId).to.be.undefined;
+  });
+
+  it('should call the wrapped function with request and context', async () => {
+    const wrapper = enrichPathInfo(mockFn);
+    await wrapper(mockRequest, mockContext);
+
+    expect(mockFn.calledOnce).to.be.true;
+    expect(mockFn.firstCall.args[0]).to.equal(mockRequest);
+    expect(mockFn.firstCall.args[1]).to.equal(mockContext);
+  });
+
+  it('should return the result from the wrapped function', async () => {
+    mockFn.resolves({ status: 201, body: 'created' });
+
+    const wrapper = enrichPathInfo(mockFn);
+    const result = await wrapper(mockRequest, mockContext);
+
+    expect(result).to.deep.equal({ status: 201, body: 'created' });
+  });
+
+  it('should handle empty pathInfo suffix', async () => {
+    mockContext.pathInfo.suffix = '';
+
+    const wrapper = enrichPathInfo(mockFn);
+    await wrapper(mockRequest, mockContext);
+
+    expect(mockContext.pathInfo.route).to.be.undefined;
+  });
+
+  it('should handle missing pathInfo suffix', async () => {
+    mockContext.pathInfo = {};
+
+    const wrapper = enrichPathInfo(mockFn);
+    await wrapper(mockRequest, mockContext);
+
+    expect(mockContext.pathInfo.route).to.be.undefined;
+  });
+
+  it('should convert method to uppercase', async () => {
+    mockRequest.method = 'get';
+
+    const wrapper = enrichPathInfo(mockFn);
+    await wrapper(mockRequest, mockContext);
+
+    expect(mockContext.pathInfo.method).to.equal('GET');
+  });
+
+  it('should handle complex route extraction', async () => {
+    mockContext.pathInfo.suffix = '/api/v1/users/123';
+
+    const wrapper = enrichPathInfo(mockFn);
+    await wrapper(mockRequest, mockContext);
+
+    expect(mockContext.pathInfo.route).to.equal('api');
+  });
+});

packages/spacecat-shared-utils/README.md

@@ -45,6 +45,40 @@ The library includes the following utility functions:
 - `hasText(str)`: Checks if the given string is not empty.
 - `dateAfterDays(number)`: Calculates the date after a specified number of days from the current date.
 
+## Log Wrapper
+
+The `logWrapper` enhances your Lambda function logs by automatically prepending `jobId` (from message) and `traceId` (from AWS X-Ray) to all log statements. This improves log traceability across distributed services.
+
+### Features
+- Automatically extracts AWS X-Ray trace ID
+- Includes jobId from message when available  
+- Enhances `context.log` directly - **no code changes needed**
+- Works seamlessly with existing log levels (info, error, debug, warn, trace, etc.)
+
+### Usage
+
+```javascript
+import { logWrapper, sqsEventAdapter } from '@adobe/spacecat-shared-utils';
+
+async function run(message, context) {
+  const { log } = context;
+
+  // Use context.log as usual - trace IDs are added automatically
+  log.info('Processing started'); 
+  // Output: [jobId=xxx] [traceId=1-xxx-xxx] Processing started
+}
+
+export const main = wrap(run)
+  .with(sqsEventAdapter)
+  .with(logWrapper)  // Add this line early in the wrapper chain
+  .with(dataAccess)
+  .with(sqs)
+  .with(secrets)
+  .with(helixStatus);
+```
+
+**Note:** The `logWrapper` enhances `context.log` directly. All existing code using `context.log` will automatically include trace IDs and job IDs in logs without any code changes.
+
 ## SQS Event Adapter
 
 The library also includes an SQS event adapter to convert an SQS record into a function parameter. This is useful when working with AWS Lambda functions that are triggered by an SQS event. Usage:
@@ -62,6 +96,21 @@ export const main = wrap(run)
   .with(helixStatus);
 ````
 
+## AWS X-Ray Integration
+
+### getTraceId()
+
+Extracts the current AWS X-Ray trace ID from the segment. Returns `null` if not in AWS Lambda or no segment is available.
+
+```javascript
+import { getTraceId } from '@adobe/spacecat-shared-utils';
+
+const traceId = getTraceId();
+// Returns: '1-5e8e8e8e-5e8e8e8e5e8e8e8e5e8e8e8e' or null
+```
+
+This function is automatically used by `logWrapper` to include trace IDs in logs.
+
 ## Testing
 
 This library includes a comprehensive test suite to ensure the reliability of the utility functions. To run the tests, use the following command:

packages/spacecat-shared-utils/src/index.d.ts

@@ -66,6 +66,36 @@ export function sqsWrapper(fn: (message: object, context: object) => Promise<Res
 export function sqsEventAdapter(fn: (message: object, context: object) => Promise<Response>):
   (request: object, context: object) => Promise<Response>;
 
+/**
+ * A higher-order function that wraps a given function and enhances logging by appending
+ * a `jobId` and `traceId` to log messages when available.
+ * @param fn - The original function to be wrapped
+ * @returns A wrapped function that enhances logging
+ */
+export function logWrapper(fn: (message: object, context: object) => Promise<Response>):
+  (message: object, context: object) => Promise<Response>;
+
+/**
+ * Instruments an AWS SDK v3 client with X-Ray tracing when running in AWS Lambda.
+ * @param client - The AWS SDK v3 client to instrument
+ * @returns The instrumented client (or original client if not in Lambda)
+ */
+export function instrumentAWSClient<T>(client: T): T;
+
+/**
+ * Extracts the trace ID from the current AWS X-Ray segment.
+ * @returns The trace ID if available, or null if not in AWS Lambda or no segment found
+ */
+export function getTraceId(): string | null;
+
+/**
+ * Adds the x-trace-id header to a headers object if a trace ID is available.
+ * @param headers - The headers object to augment
+ * @param context - The context object that may contain traceId
+ * @returns The headers object with x-trace-id added if available
+ */
+export function addTraceIdHeader(headers?: Record<string, string>, context?: object): Record<string, string>;
+
 /**
  * Prepends 'https://' schema to the URL if it's not already present.
  * @param url - The URL to modify.

packages/spacecat-shared-utils/src/index.js

@@ -52,7 +52,7 @@ export { sqsWrapper } from './sqs.js';
 export { sqsEventAdapter } from './sqs.js';
 
 export { logWrapper } from './log-wrapper.js';
-export { instrumentAWSClient } from './xray.js';
+export { instrumentAWSClient, getTraceId, addTraceIdHeader } from './xray.js';
 
 export {
   composeBaseURL,

packages/spacecat-shared-utils/src/log-wrapper.js

@@ -10,47 +10,63 @@
  * governing permissions and limitations under the License.
  */
 
+import { getTraceId } from './xray.js';
+
 /**
  * A higher-order function that wraps a given function and enhances logging by appending
- * a `jobId` to log messages when available. This improves traceability of logs associated
- * with specific jobs or processes.
+ * a `jobId` and `traceId` to log messages when available. This improves traceability of logs
+ * associated with specific jobs or processes.
  *
  * The wrapper checks if a `log` object exists in the `context` and whether the `message`
- * contains a `jobId`. If found, log methods (e.g., `info`, `error`, etc.) will prepend the
- * `jobId` to all log statements where `context.contextualLog` is used. If no `jobId` is found,
- * logging will remain unchanged.
+ * contains a `jobId`. It also extracts the AWS X-Ray trace ID if available. If found, log
+ * methods (e.g., `info`, `error`, etc.) will prepend the `jobId` and/or `traceId` to all log
+ * statements. All existing code using `context.log` will automatically include these markers.
  *
  * @param {function} fn - The original function to be wrapped, called with the provided
  * message and context after logging enhancement.
  * @returns {function(object, object): Promise<Response>} - A wrapped function that enhances
  * logging and returns the result of the original function.
  *
- * `context.contextualLog` will include logging methods with `jobId` prefixed, or fall back
- * to the existing `log` object if no `jobId` is provided.
+ * `context.log` will be enhanced in place to include `jobId` and/or `traceId` prefixed to all
+ * log messages. No code changes needed - existing `context.log` calls work automatically.
  */
 export function logWrapper(fn) {
   return async (message, context) => {
     const { log } = context;
 
     if (log && !context.contextualLog) {
+      const markers = [];
+
+      // Extract jobId from message if available
       if (typeof message === 'object' && message !== null && 'jobId' in message) {
         const { jobId } = message;
-        const jobIdMarker = `[jobId=${jobId}]`;
+        markers.push(`[jobId=${jobId}]`);
+      }
+
+      // Extract traceId from AWS X-Ray
+      const traceId = getTraceId();
+      if (traceId) {
+        markers.push(`[traceId=${traceId}]`);
+      }
+
+      // If we have markers, enhance the log object directly
+      if (markers.length > 0) {
+        const markerString = markers.join(' ');
 
         // Define log levels
         const logLevels = ['info', 'error', 'debug', 'warn', 'trace', 'verbose', 'silly', 'fatal'];
 
-        // Enhance the log object to include jobId in all log statements
-        context.contextualLog = logLevels.reduce((accumulator, level) => {
+        // Enhance context.log directly to include markers in all log statements
+        context.log = logLevels.reduce((accumulator, level) => {
           if (typeof log[level] === 'function') {
-            accumulator[level] = (...args) => log[level](jobIdMarker, ...args);
+            accumulator[level] = (...args) => log[level](markerString, ...args);
           }
           return accumulator;
         }, {});
-      } else {
-        log.debug('No jobId found in the provided message. Log entries will be recorded without a jobId.');
-        context.contextualLog = log;
       }
+
+      // Mark that we've processed this context
+      context.contextualLog = context.log;
     }
 
     return fn(message, context);