feat: emit X-Ray subsegment on every lambda:// invoke

package.json

@@ -58,6 +58,7 @@
     "@aws-sdk/signature-v4": "^3.110.0",
     "@aws-sdk/url-parser": "^3.357.0",
     "@types/aws-lambda": "^8.10.101",
+    "aws-xray-sdk-core": "^3.12.0",
     "axios": "^1.6.0",
     "lodash": "^4.17.21",
     "nearley": "2",
@@ -66,5 +67,8 @@
   },
   "publishConfig": {
     "access": "public"
+  },
+  "resolutions": {
+    "@aws-sdk/types": "3.357.0"
   }
 }

src/adapters/helpers/tracing.ts

@@ -0,0 +1,58 @@
+import { Lambda } from '@aws-sdk/client-lambda';
+import { captureAWSv3Client, setContextMissingStrategy } from 'aws-xray-sdk-core';
+
+/**
+ * Whether X-Ray tracing is active in the current runtime.
+ *
+ * AWS injects `_X_AMZN_TRACE_ID` into any Lambda execution environment that has
+ * active tracing, so its presence is a reliable signal that there is a live
+ * segment to attach to. Consumers can also force-enable instrumentation via
+ * `ALPHA_XRAY_TRACING=true` (e.g. for a non-Lambda runtime that manages its own
+ * segments).
+ *
+ * Outside of these cases (local dev, tests, untraced runtimes) we deliberately
+ * leave alpha's Lambda client untouched: alpha is a published library used
+ * everywhere, and wrapping an untraced client would emit context-missing errors
+ * on every invoke. The default behavior for an untraced caller must be
+ * unchanged.
+ */
+const isXRayTracingActive = (): boolean =>
+  Boolean(process.env._X_AMZN_TRACE_ID) || process.env.ALPHA_XRAY_TRACING === 'true';
+
+// Only soften the context-missing strategy once, and never more than once.
+let contextMissingConfigured = false;
+
+/**
+ * Wrap alpha's Lambda client with X-Ray so that each `lambda://` invoke emits a
+ * subsegment naming the downstream Lambda (the target service), with trace
+ * header propagation. This lights up internal service -> service edges in the
+ * X-Ray service graph.
+ *
+ * `captureAWSv3Client` adds middleware to and returns the *same* client
+ * instance, so `lambda.destroy()` and the `config.Lambda` injection escape
+ * hatch keep working transparently against the returned client.
+ *
+ * Safety for untraced consumers is the priority here:
+ *  - We only wrap when X-Ray is actually active (see `isXRayTracingActive`).
+ *  - If `AWS_XRAY_CONTEXT_MISSING` is unset, we set the strategy to `LOG_ERROR`
+ *    so a missing segment logs rather than throws. Consumers that need a
+ *    different strategy should set that env var in the runtime environment.
+ *  - Any failure while instrumenting falls back to the original, unwrapped
+ *    client so tracing can never break a real invoke.
+ */
+export const captureLambdaClient = (lambda: Lambda): Lambda => {
+  if (!isXRayTracingActive()) {
+    return lambda;
+  }
+
+  try {
+    if (!contextMissingConfigured && !process.env.AWS_XRAY_CONTEXT_MISSING) {
+      setContextMissingStrategy('LOG_ERROR');
+      contextMissingConfigured = true;
+    }
+    return captureAWSv3Client(lambda);
+  } catch {
+    // Never let tracing instrumentation break a real invoke.
+    return lambda;
+  }
+};

src/adapters/lambda-invocation.ts

@@ -6,6 +6,7 @@ import { lambdaEvent } from './helpers/lambdaEvent';
 import { lambdaResponse, Payload } from './helpers/lambdaResponse';
 import { parseLambdaUrl, isAbsoluteURL } from '../utils/url';
 import { RequestError } from './helpers/requestError';
+import { captureLambdaClient } from './helpers/tracing';
 import { InternalAlphaRequestConfig, AlphaAdapter } from '../types';
 import { Alpha } from '../alpha';
 import { AbortController } from '@aws-sdk/abort-controller';
@@ -32,7 +33,11 @@ const lambdaInvocationAdapter: AlphaAdapter = async (config) => {
     });
   }
 
-  const lambda = new LambdaClass(lambdaOptions);
+  // Wrap the client with X-Ray (safely, no-op when untraced) so that every
+  // `lambda://` invoke emits a subsegment naming the downstream service. This
+  // also covers the `config.Lambda` injection escape hatch, so all internal
+  // service -> service traffic is traced through this single chokepoint.
+  const lambda = captureLambdaClient(new LambdaClass(lambdaOptions));
   if (config.baseURL && !isAbsoluteURL(config.url as string)) {
     config.url = `${config.baseURL}${config.url}`;
   }

test/lambda-invocation.test.ts

@@ -12,6 +12,16 @@ import {
   lambdaResponse,
   Payload,
 } from '../src/adapters/helpers/lambdaResponse';
+import { captureAWSv3Client } from 'aws-xray-sdk-core';
+
+jest.mock('aws-xray-sdk-core', () => ({
+  captureAWSv3Client: jest.fn((client) => client),
+  setContextMissingStrategy: jest.fn(),
+}));
+
+const mockCaptureAWSv3Client = captureAWSv3Client as jest.MockedFunction<
+  typeof captureAWSv3Client
+>;
 
 const mockLambda = mockClient(Lambda);
 const FakeLambda = jest.fn() as jest.MockedClass<typeof Lambda>;
@@ -608,3 +618,129 @@ test('lambdaRegion config option is provided to the Lambda client', async () =>
 
   expect(FakeLambda).toHaveBeenCalledWith({ region: 'ap-southeast-2' });
 });
+
+describe('X-Ray tracing', () => {
+  const originalTraceId = process.env._X_AMZN_TRACE_ID;
+  const originalOptIn = process.env.ALPHA_XRAY_TRACING;
+
+  beforeEach(() => {
+    // `resetMocks: true` wipes implementations between tests, so (re)install
+    // the pass-through behavior captureAWSv3Client has by default.
+    mockCaptureAWSv3Client.mockImplementation((client) => client);
+    delete process.env._X_AMZN_TRACE_ID;
+    delete process.env.ALPHA_XRAY_TRACING;
+  });
+
+  afterAll(() => {
+    if (originalTraceId === undefined) {
+      delete process.env._X_AMZN_TRACE_ID;
+    } else {
+      process.env._X_AMZN_TRACE_ID = originalTraceId;
+    }
+    if (originalOptIn === undefined) {
+      delete process.env.ALPHA_XRAY_TRACING;
+    } else {
+      process.env.ALPHA_XRAY_TRACING = originalOptIn;
+    }
+  });
+
+  test('a lambda:// invoke wraps the Lambda client when ALPHA_XRAY_TRACING is enabled', async () => {
+    process.env.ALPHA_XRAY_TRACING = 'true';
+    createResponse(mockLambda, {
+      StatusCode: 200,
+      Payload: {
+        body: 'hello!',
+        headers: { 'test-header': 'some value' },
+        statusCode: 200,
+      },
+    });
+
+    const response = await ctx.alpha.get('/some/path');
+
+    expect(response.data).toBe('hello!');
+    expect(response.status).toBe(200);
+    expect(mockCaptureAWSv3Client).toHaveBeenCalledTimes(1);
+    expect(mockCaptureAWSv3Client).toHaveBeenCalledWith(expect.any(Lambda));
+  });
+
+  test('a lambda:// invoke wraps the Lambda client for tracing when X-Ray is active', async () => {
+    process.env._X_AMZN_TRACE_ID = 'Root=1-5e1b4151-5ac6c58f5b3e6f6f00000000';
+    createResponse(mockLambda, {
+      StatusCode: 200,
+      Payload: {
+        body: 'hello!',
+        headers: { 'test-header': 'some value' },
+        statusCode: 200,
+      },
+    });
+
+    const response = await ctx.alpha.get('/some/path');
+
+    // The downstream invoke still returns the expected payload...
+    expect(response.data).toBe('hello!');
+    expect(response.status).toBe(200);
+
+    // ...and the Lambda client used for the invoke was wrapped by X-Ray.
+    expect(mockCaptureAWSv3Client).toHaveBeenCalledTimes(1);
+    expect(mockCaptureAWSv3Client).toHaveBeenCalledWith(expect.any(Lambda));
+  });
+
+  test('an injected config.Lambda client is still wrapped for tracing', async () => {
+    process.env._X_AMZN_TRACE_ID = 'Root=1-5e1b4151-5ac6c58f5b3e6f6f00000000';
+    createResponse(mockLambda, {
+      StatusCode: 200,
+      Payload: {
+        body: 'test',
+        statusCode: 200,
+      },
+    });
+
+    const response = await ctx.alpha.get('/test', { Lambda: FakeLambda });
+
+    expect(response.data).toBe('test');
+    expect(FakeLambda).toHaveBeenCalledTimes(1);
+    // The instance produced by the injected class is what gets traced.
+    expect(mockCaptureAWSv3Client).toHaveBeenCalledTimes(1);
+    expect(mockCaptureAWSv3Client).toHaveBeenCalledWith(expect.any(Lambda));
+  });
+
+  test('the untraced path is safe: the client is not wrapped and the invoke succeeds', async () => {
+    // No _X_AMZN_TRACE_ID and no opt-in => tracing must be a no-op.
+    createResponse(mockLambda, {
+      StatusCode: 200,
+      Payload: {
+        body: 'hello!',
+        headers: { 'test-header': 'some value' },
+        statusCode: 200,
+      },
+    });
+
+    const response = await ctx.alpha.get('/some/path');
+
+    expect(response.data).toBe('hello!');
+    expect(response.status).toBe(200);
+    expect(mockCaptureAWSv3Client).not.toHaveBeenCalled();
+  });
+
+  test('instrumentation failures never break the invoke (falls back to the unwrapped client)', async () => {
+    process.env._X_AMZN_TRACE_ID = 'Root=1-5e1b4151-5ac6c58f5b3e6f6f00000000';
+    mockCaptureAWSv3Client.mockImplementation(() => {
+      throw new Error('boom');
+    });
+    createResponse(mockLambda, {
+      StatusCode: 200,
+      Payload: {
+        body: 'hello!',
+        statusCode: 200,
+      },
+    });
+
+    const response = await ctx.alpha.get('/some/path');
+
+    // The wrap was attempted but threw; the original client still serviced the
+    // invoke, so consumers are never broken by tracing.
+    expect(mockCaptureAWSv3Client).toHaveBeenCalledTimes(1);
+    expect(response.data).toBe('hello!');
+    expect(response.status).toBe(200);
+  });
+});