Skip to content

Commit ed1fb61

Browse files
LLMO-1023: add trace ID propagation for SQS messages and HTTP headers (#1097)
## Overview Adds end-to-end trace ID propagation across SpaceCat services, enabling distributed tracing for debugging and monitoring. ## Changes ### 1. Log Wrapper Enhancement (`spacecat-shared-utils`) - Enhanced `logWrapper` to automatically include `traceId` in log messages - Extracts trace ID from AWS X-Ray or context for correlation ### 2. SQS Message Propagation (`spacecat-shared-utils`) - `sendMessage()` automatically includes `traceId` in message payload - `sqsEventAdapter()` extracts `traceId` from incoming messages and stores in `context.traceId` ### 3. HTTP Header Support (`spacecat-shared-http-utils`) - `enrichPathInfo()` extracts `x-trace-id` header from incoming HTTP requests - Stores trace ID in `context.traceId` for propagation ### 4. New Utility Function (`spacecat-shared-utils`) - `addTraceIdHeader(headers, context)` - Helper to add `x-trace-id` to outgoing HTTP requests - Exported from `@adobe/spacecat-shared-utils`
1 parent 3de80b1 commit ed1fb61

File tree

13 files changed

+780
-42
lines changed

13 files changed

+780
-42
lines changed

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,14 +13,23 @@
1313
export function enrichPathInfo(fn) { // export for testing
1414
return async (request, context) => {
1515
const [, route] = context?.pathInfo?.suffix?.split(/\/+/) || [];
16+
const headers = request.headers.plain();
17+
1618
context.pathInfo = {
1719
...context.pathInfo,
1820
...{
1921
method: request.method.toUpperCase(),
20-
headers: request.headers.plain(),
22+
headers,
2123
route,
2224
},
2325
};
26+
27+
// Extract and store traceId from x-trace-id header if present
28+
const traceIdHeader = headers['x-trace-id'];
29+
if (traceIdHeader) {
30+
context.traceId = traceIdHeader;
31+
}
32+
2433
return fn(request, context);
2534
};
2635
}
Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
/*
2+
* Copyright 2025 Adobe. All rights reserved.
3+
* This file is licensed to you under the Apache License, Version 2.0 (the "License");
4+
* you may not use this file except in compliance with the License. You may obtain a copy
5+
* of the License at http://www.apache.org/licenses/LICENSE-2.0
6+
*
7+
* Unless required by applicable law or agreed to in writing, software distributed under
8+
* the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS
9+
* OF ANY KIND, either express or implied. See the License for the specific language
10+
* governing permissions and limitations under the License.
11+
*/
12+
/* eslint-env mocha */
13+
import { expect } from 'chai';
14+
import sinon from 'sinon';
15+
16+
import { enrichPathInfo } from '../src/enrich-path-info-wrapper.js';
17+
18+
describe('enrichPathInfo', () => {
19+
let mockRequest;
20+
let mockContext;
21+
let mockFn;
22+
23+
beforeEach(() => {
24+
mockFn = sinon.stub().resolves({ status: 200 });
25+
26+
mockRequest = {
27+
method: 'POST',
28+
headers: {
29+
plain: () => ({
30+
'content-type': 'application/json',
31+
'user-agent': 'test-agent',
32+
}),
33+
},
34+
};
35+
36+
mockContext = {
37+
pathInfo: {
38+
suffix: '/api/test',
39+
},
40+
};
41+
});
42+
43+
afterEach(() => {
44+
sinon.restore();
45+
});
46+
47+
it('should enrich context with pathInfo including method, headers, and route', async () => {
48+
const wrapper = enrichPathInfo(mockFn);
49+
await wrapper(mockRequest, mockContext);
50+
51+
expect(mockContext.pathInfo).to.deep.include({
52+
method: 'POST',
53+
route: 'api',
54+
});
55+
expect(mockContext.pathInfo.headers).to.deep.equal({
56+
'content-type': 'application/json',
57+
'user-agent': 'test-agent',
58+
});
59+
});
60+
61+
it('should extract traceId from x-trace-id header and store in context', async () => {
62+
mockRequest.headers.plain = () => ({
63+
'content-type': 'application/json',
64+
'x-trace-id': '1-5e8e8e8e-5e8e8e8e5e8e8e8e5e8e8e8e',
65+
});
66+
67+
const wrapper = enrichPathInfo(mockFn);
68+
await wrapper(mockRequest, mockContext);
69+
70+
expect(mockContext.traceId).to.equal('1-5e8e8e8e-5e8e8e8e5e8e8e8e5e8e8e8e');
71+
});
72+
73+
it('should not set traceId in context when x-trace-id header is missing', async () => {
74+
mockRequest.headers.plain = () => ({
75+
'content-type': 'application/json',
76+
});
77+
78+
const wrapper = enrichPathInfo(mockFn);
79+
await wrapper(mockRequest, mockContext);
80+
81+
expect(mockContext.traceId).to.be.undefined;
82+
});
83+
84+
it('should handle case-sensitive x-trace-id header', async () => {
85+
mockRequest.headers.plain = () => ({
86+
'content-type': 'application/json',
87+
'X-Trace-Id': '1-different-case',
88+
});
89+
90+
const wrapper = enrichPathInfo(mockFn);
91+
await wrapper(mockRequest, mockContext);
92+
93+
// Header keys should be lowercase
94+
expect(mockContext.traceId).to.be.undefined;
95+
});
96+
97+
it('should call the wrapped function with request and context', async () => {
98+
const wrapper = enrichPathInfo(mockFn);
99+
await wrapper(mockRequest, mockContext);
100+
101+
expect(mockFn.calledOnce).to.be.true;
102+
expect(mockFn.firstCall.args[0]).to.equal(mockRequest);
103+
expect(mockFn.firstCall.args[1]).to.equal(mockContext);
104+
});
105+
106+
it('should return the result from the wrapped function', async () => {
107+
mockFn.resolves({ status: 201, body: 'created' });
108+
109+
const wrapper = enrichPathInfo(mockFn);
110+
const result = await wrapper(mockRequest, mockContext);
111+
112+
expect(result).to.deep.equal({ status: 201, body: 'created' });
113+
});
114+
115+
it('should handle empty pathInfo suffix', async () => {
116+
mockContext.pathInfo.suffix = '';
117+
118+
const wrapper = enrichPathInfo(mockFn);
119+
await wrapper(mockRequest, mockContext);
120+
121+
expect(mockContext.pathInfo.route).to.be.undefined;
122+
});
123+
124+
it('should handle missing pathInfo suffix', async () => {
125+
mockContext.pathInfo = {};
126+
127+
const wrapper = enrichPathInfo(mockFn);
128+
await wrapper(mockRequest, mockContext);
129+
130+
expect(mockContext.pathInfo.route).to.be.undefined;
131+
});
132+
133+
it('should convert method to uppercase', async () => {
134+
mockRequest.method = 'get';
135+
136+
const wrapper = enrichPathInfo(mockFn);
137+
await wrapper(mockRequest, mockContext);
138+
139+
expect(mockContext.pathInfo.method).to.equal('GET');
140+
});
141+
142+
it('should handle complex route extraction', async () => {
143+
mockContext.pathInfo.suffix = '/api/v1/users/123';
144+
145+
const wrapper = enrichPathInfo(mockFn);
146+
await wrapper(mockRequest, mockContext);
147+
148+
expect(mockContext.pathInfo.route).to.equal('api');
149+
});
150+
});

packages/spacecat-shared-utils/README.md

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,40 @@ The library includes the following utility functions:
4545
- `hasText(str)`: Checks if the given string is not empty.
4646
- `dateAfterDays(number)`: Calculates the date after a specified number of days from the current date.
4747

48+
## Log Wrapper
49+
50+
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.
51+
52+
### Features
53+
- Automatically extracts AWS X-Ray trace ID
54+
- Includes jobId from message when available
55+
- Enhances `context.log` directly - **no code changes needed**
56+
- Works seamlessly with existing log levels (info, error, debug, warn, trace, etc.)
57+
58+
### Usage
59+
60+
```javascript
61+
import { logWrapper, sqsEventAdapter } from '@adobe/spacecat-shared-utils';
62+
63+
async function run(message, context) {
64+
const { log } = context;
65+
66+
// Use context.log as usual - trace IDs are added automatically
67+
log.info('Processing started');
68+
// Output: [jobId=xxx] [traceId=1-xxx-xxx] Processing started
69+
}
70+
71+
export const main = wrap(run)
72+
.with(sqsEventAdapter)
73+
.with(logWrapper) // Add this line early in the wrapper chain
74+
.with(dataAccess)
75+
.with(sqs)
76+
.with(secrets)
77+
.with(helixStatus);
78+
```
79+
80+
**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.
81+
4882
## SQS Event Adapter
4983

5084
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)
6296
.with(helixStatus);
6397
````
6498

99+
## AWS X-Ray Integration
100+
101+
### getTraceId()
102+
103+
Extracts the current AWS X-Ray trace ID from the segment. Returns `null` if not in AWS Lambda or no segment is available.
104+
105+
```javascript
106+
import { getTraceId } from '@adobe/spacecat-shared-utils';
107+
108+
const traceId = getTraceId();
109+
// Returns: '1-5e8e8e8e-5e8e8e8e5e8e8e8e5e8e8e8e' or null
110+
```
111+
112+
This function is automatically used by `logWrapper` to include trace IDs in logs.
113+
65114
## Testing
66115

67116
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

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,36 @@ export function sqsWrapper(fn: (message: object, context: object) => Promise<Res
6666
export function sqsEventAdapter(fn: (message: object, context: object) => Promise<Response>):
6767
(request: object, context: object) => Promise<Response>;
6868

69+
/**
70+
* A higher-order function that wraps a given function and enhances logging by appending
71+
* a `jobId` and `traceId` to log messages when available.
72+
* @param fn - The original function to be wrapped
73+
* @returns A wrapped function that enhances logging
74+
*/
75+
export function logWrapper(fn: (message: object, context: object) => Promise<Response>):
76+
(message: object, context: object) => Promise<Response>;
77+
78+
/**
79+
* Instruments an AWS SDK v3 client with X-Ray tracing when running in AWS Lambda.
80+
* @param client - The AWS SDK v3 client to instrument
81+
* @returns The instrumented client (or original client if not in Lambda)
82+
*/
83+
export function instrumentAWSClient<T>(client: T): T;
84+
85+
/**
86+
* Extracts the trace ID from the current AWS X-Ray segment.
87+
* @returns The trace ID if available, or null if not in AWS Lambda or no segment found
88+
*/
89+
export function getTraceId(): string | null;
90+
91+
/**
92+
* Adds the x-trace-id header to a headers object if a trace ID is available.
93+
* @param headers - The headers object to augment
94+
* @param context - The context object that may contain traceId
95+
* @returns The headers object with x-trace-id added if available
96+
*/
97+
export function addTraceIdHeader(headers?: Record<string, string>, context?: object): Record<string, string>;
98+
6999
/**
70100
* Prepends 'https://' schema to the URL if it's not already present.
71101
* @param url - The URL to modify.

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ export { sqsWrapper } from './sqs.js';
5252
export { sqsEventAdapter } from './sqs.js';
5353

5454
export { logWrapper } from './log-wrapper.js';
55-
export { instrumentAWSClient } from './xray.js';
55+
export { instrumentAWSClient, getTraceId, addTraceIdHeader } from './xray.js';
5656

5757
export {
5858
composeBaseURL,

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

Lines changed: 30 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -10,47 +10,63 @@
1010
* governing permissions and limitations under the License.
1111
*/
1212

13+
import { getTraceId } from './xray.js';
14+
1315
/**
1416
* A higher-order function that wraps a given function and enhances logging by appending
15-
* a `jobId` to log messages when available. This improves traceability of logs associated
16-
* with specific jobs or processes.
17+
* a `jobId` and `traceId` to log messages when available. This improves traceability of logs
18+
* associated with specific jobs or processes.
1719
*
1820
* The wrapper checks if a `log` object exists in the `context` and whether the `message`
19-
* contains a `jobId`. If found, log methods (e.g., `info`, `error`, etc.) will prepend the
20-
* `jobId` to all log statements where `context.contextualLog` is used. If no `jobId` is found,
21-
* logging will remain unchanged.
21+
* contains a `jobId`. It also extracts the AWS X-Ray trace ID if available. If found, log
22+
* methods (e.g., `info`, `error`, etc.) will prepend the `jobId` and/or `traceId` to all log
23+
* statements. All existing code using `context.log` will automatically include these markers.
2224
*
2325
* @param {function} fn - The original function to be wrapped, called with the provided
2426
* message and context after logging enhancement.
2527
* @returns {function(object, object): Promise<Response>} - A wrapped function that enhances
2628
* logging and returns the result of the original function.
2729
*
28-
* `context.contextualLog` will include logging methods with `jobId` prefixed, or fall back
29-
* to the existing `log` object if no `jobId` is provided.
30+
* `context.log` will be enhanced in place to include `jobId` and/or `traceId` prefixed to all
31+
* log messages. No code changes needed - existing `context.log` calls work automatically.
3032
*/
3133
export function logWrapper(fn) {
3234
return async (message, context) => {
3335
const { log } = context;
3436

3537
if (log && !context.contextualLog) {
38+
const markers = [];
39+
40+
// Extract jobId from message if available
3641
if (typeof message === 'object' && message !== null && 'jobId' in message) {
3742
const { jobId } = message;
38-
const jobIdMarker = `[jobId=${jobId}]`;
43+
markers.push(`[jobId=${jobId}]`);
44+
}
45+
46+
// Extract traceId from AWS X-Ray
47+
const traceId = getTraceId();
48+
if (traceId) {
49+
markers.push(`[traceId=${traceId}]`);
50+
}
51+
52+
// If we have markers, enhance the log object directly
53+
if (markers.length > 0) {
54+
const markerString = markers.join(' ');
3955

4056
// Define log levels
4157
const logLevels = ['info', 'error', 'debug', 'warn', 'trace', 'verbose', 'silly', 'fatal'];
4258

43-
// Enhance the log object to include jobId in all log statements
44-
context.contextualLog = logLevels.reduce((accumulator, level) => {
59+
// Enhance context.log directly to include markers in all log statements
60+
context.log = logLevels.reduce((accumulator, level) => {
4561
if (typeof log[level] === 'function') {
46-
accumulator[level] = (...args) => log[level](jobIdMarker, ...args);
62+
accumulator[level] = (...args) => log[level](markerString, ...args);
4763
}
4864
return accumulator;
4965
}, {});
50-
} else {
51-
log.debug('No jobId found in the provided message. Log entries will be recorded without a jobId.');
52-
context.contextualLog = log;
5366
}
67+
68+
// Mark that we've processed this context
69+
context.contextualLog = context.log;
5470
}
5571

5672
return fn(message, context);

0 commit comments

Comments
 (0)