Check for undefined XRay Trace Id in AwsSdk Instrumentation

aws-distro-opentelemetry-node-autoinstrumentation/src/patches/instrumentation-patch.ts

@@ -311,9 +311,12 @@ function patchAwsSdkInstrumentation(instrumentation: Instrumentation): void {
             // Need to set capitalized version of the trace id to ensure that the Recursion Detection Middleware
             // of aws-sdk-js-v3 will detect the propagated X-Ray Context
             // See: https://github.com/aws/aws-sdk-js-v3/blob/v3.768.0/packages/middleware-recursion-detection/src/index.ts#L13
-            middlewareArgs.request.headers[AWSXRAY_TRACE_ID_HEADER_CAPITALIZED] =
-              middlewareArgs.request.headers[AWSXRAY_TRACE_ID_HEADER];
-            delete middlewareArgs.request.headers[AWSXRAY_TRACE_ID_HEADER];
+            const xrayTraceId = middlewareArgs.request.headers[AWSXRAY_TRACE_ID_HEADER];
+
+            if (xrayTraceId) {
+              middlewareArgs.request.headers[AWSXRAY_TRACE_ID_HEADER_CAPITALIZED] = xrayTraceId;
+              delete middlewareArgs.request.headers[AWSXRAY_TRACE_ID_HEADER];
+            }
             const result = await next(middlewareArgs);
             return result;
           },

aws-distro-opentelemetry-node-autoinstrumentation/test/patches/instrumentation-patch.test.ts

@@ -535,34 +535,67 @@ describe('InstrumentationPatchTest', () => {
     let lambda: Lambda;
     const region = 'us-east-1';
 
-    it('overridden _getV3SmithyClientSendPatch updates MiddlewareStack', async () => {
-      const mockedMiddlewareStackInternal: any = [];
-      const mockedMiddlewareStack = {
-        add: (arg1: any, arg2: any) => mockedMiddlewareStackInternal.push([arg1, arg2]),
-      };
-      const send = extractAwsSdkInstrumentation(PATCHED_INSTRUMENTATIONS)
-        ['_getV3SmithyClientSendPatch']((...args: unknown[]) => Promise.resolve())
-        .bind({ middlewareStack: mockedMiddlewareStack });
-      sinon
-        .stub(AWSXRayPropagator.prototype, 'inject')
-        .callsFake((context: OtelContext, carrier: unknown, setter: TextMapSetter) => {
-          (carrier as any)['isCarrierModified'] = 'carrierIsModified';
-        });
+    describe('overridden _getV3SmithyClientSendPatch updates MiddlewareStack', async () => {
+      let mockedMiddlewareStackInternal: any;
+      let mockedMiddlewareStack;
+      let send;
+
+      beforeEach(async () => {
+        // Clear environment variables before each test
+        mockedMiddlewareStackInternal = [];
+        mockedMiddlewareStack = {
+          add: (arg1: any, arg2: any) => mockedMiddlewareStackInternal.push([arg1, arg2]),
+        };
+        send = extractAwsSdkInstrumentation(PATCHED_INSTRUMENTATIONS)
+          ['_getV3SmithyClientSendPatch']((...args: unknown[]) => Promise.resolve())
+          .bind({ middlewareStack: mockedMiddlewareStack });
+
+        sinon
+          .stub(AWSXRayPropagator.prototype, 'inject')
+          .callsFake((context: OtelContext, carrier: unknown, setter: TextMapSetter) => {
+            (carrier as any)['isCarrierModified'] = 'carrierIsModified';
+          });
+
+        // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+        // @ts-ignore
+        await send({}, null);
+      });
 
-      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-      // @ts-ignore
-      await send({}, null);
+      it('Injecting with existing X-Ray header', async () => {
+        const existingHeader = 'test-trace-header';
+        const middlewareArgsWithHeader: any = {
+          request: {
+            headers: { [AWSXRAY_TRACE_ID_HEADER]: existingHeader },
+          },
+        };
+        await mockedMiddlewareStackInternal[0][0]((arg: any) => Promise.resolve(), null)(middlewareArgsWithHeader);
 
-      const middlewareArgs: any = {
-        request: {
-          headers: {},
-        },
-      };
-      await mockedMiddlewareStackInternal[0][0]((arg: any) => Promise.resolve(), null)(middlewareArgs);
+        sinon.restore();
+
+        expect(middlewareArgsWithHeader.request.headers['isCarrierModified']).toEqual('carrierIsModified');
+        expect(middlewareArgsWithHeader.request.headers).not.toHaveProperty(AWSXRAY_TRACE_ID_HEADER);
+        expect(middlewareArgsWithHeader.request.headers).toHaveProperty('X-Amzn-Trace-Id');
+        expect(middlewareArgsWithHeader.request.headers['X-Amzn-Trace-Id']).toEqual(existingHeader);
+
+        expect(mockedMiddlewareStackInternal[0][1].name).toEqual('_adotInjectXrayContextMiddleware');
+      });
 
-      sinon.restore();
-      expect(middlewareArgs.request.headers['isCarrierModified']).toEqual('carrierIsModified');
-      expect(mockedMiddlewareStackInternal[0][1].name).toEqual('_adotInjectXrayContextMiddleware');
+      it('Injecting without existing X-Ray header', async () => {
+        const middlewareArgsNoHeader: any = {
+          request: {
+            headers: {},
+          },
+        };
+
+        await mockedMiddlewareStackInternal[0][0]((arg: any) => Promise.resolve(), null)(middlewareArgsNoHeader);
+
+        sinon.restore();
+
+        expect(middlewareArgsNoHeader.request.headers['isCarrierModified']).toEqual('carrierIsModified');
+        expect(middlewareArgsNoHeader.request.headers).not.toHaveProperty('X-Amzn-Trace-Id');
+
+        expect(mockedMiddlewareStackInternal[0][1].name).toEqual('_adotInjectXrayContextMiddleware');
+      });
     });
 
     it('injects trace context header into request via propagator', async () => {