feat(core/client): update trace header conditions

Author: kuhe

packages-internal/core/src/submodules/client/middleware-recursion-detection/middleware-recursion-detection.integ.spec.ts

@@ -128,6 +128,91 @@ describe("middleware-recursion-detection", () => {
       expect.hasAssertions();
     });
 
+    it("should NOT overwrite existing traceparent or add tracestate/baggage from InvokeStore when traceparent is already set", async () => {
+      const existingTraceparent = "00-abcdef1234567890abcdef1234567890-1234567890abcdef-01";
+      const mockTraceparent = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";
+      const mockTracestate = "congo=t61rcWkgMzE";
+      const mockBaggage = "userId=alice,serverNode=DF%2028";
+      vi.spyOn(InvokeStore, "getInstanceAsync").mockResolvedValue({
+        getXRayTraceId: () => undefined,
+        getTraceparent: () => mockTraceparent,
+        getTracestate: () => mockTracestate,
+        getBaggage: () => mockBaggage,
+      } as any);
+
+      const client = new Lambda({
+        region: "us-west-2",
+      });
+
+      // Add middleware that sets traceparent before recursion detection runs.
+      // recursionDetectionMiddleware runs at build step with low priority,
+      // so we add at build step with high priority to run first.
+      client.middlewareStack.add(
+        (next) => async (args: any) => {
+          args.request.headers["traceparent"] = existingTraceparent;
+          return next(args);
+        },
+        { step: "build", priority: "high", name: "testSetTraceparent", override: true }
+      );
+
+      requireRequestsFrom(client).toMatch({
+        headers: {
+          traceparent: new RegExp(`^${existingTraceparent.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}$`),
+          tracestate: (value: any) => expect(value).toBeUndefined(),
+          baggage: (value: any) => expect(value).toBeUndefined(),
+        },
+      });
+
+      await client.invoke({
+        FunctionName: "my-function",
+      });
+
+      expect.hasAssertions();
+    });
+
+    it("should preserve existing traceparent, tracestate, and baggage without overwriting from InvokeStore", async () => {
+      const existingTraceparent = "00-abcdef1234567890abcdef1234567890-1234567890abcdef-01";
+      const existingTracestate = "vendor1=opaqueValue1";
+      const existingBaggage = "key1=value1";
+      const mockTraceparent = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";
+      const mockTracestate = "congo=t61rcWkgMzE";
+      const mockBaggage = "userId=alice,serverNode=DF%2028";
+      vi.spyOn(InvokeStore, "getInstanceAsync").mockResolvedValue({
+        getXRayTraceId: () => undefined,
+        getTraceparent: () => mockTraceparent,
+        getTracestate: () => mockTracestate,
+        getBaggage: () => mockBaggage,
+      } as any);
+
+      const client = new Lambda({
+        region: "us-west-2",
+      });
+
+      client.middlewareStack.add(
+        (next) => async (args: any) => {
+          args.request.headers["traceparent"] = existingTraceparent;
+          args.request.headers["tracestate"] = existingTracestate;
+          args.request.headers["baggage"] = existingBaggage;
+          return next(args);
+        },
+        { step: "build", priority: "high", name: "testSetTraceHeaders", override: true }
+      );
+
+      requireRequestsFrom(client).toMatch({
+        headers: {
+          traceparent: new RegExp(`^${existingTraceparent.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}$`),
+          tracestate: new RegExp(`^${existingTracestate.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}$`),
+          baggage: new RegExp(`^${existingBaggage.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}$`),
+        },
+      });
+
+      await client.invoke({
+        FunctionName: "my-function",
+      });
+
+      expect.hasAssertions();
+    });
+
     it("should NOT add W3C headers when InvokeStore has no traceparent", async () => {
       vi.spyOn(InvokeStore, "getInstanceAsync").mockResolvedValue({
         getXRayTraceId: () => undefined,

packages-internal/core/src/submodules/client/middleware-recursion-detection/recursionDetectionMiddleware.spec.ts

@@ -256,7 +256,7 @@ describe(recursionDetectionMiddleware.name, () => {
       expect(request.headers["baggage"]).toBeUndefined();
     });
 
-    it("should preserve existing traceparent header and add tracestate/baggage from InvokeStore", async () => {
+    it("should preserve existing traceparent header and NOT add tracestate/baggage from InvokeStore", async () => {
       const existingTraceparent = "00-abcdef1234567890abcdef1234567890-1234567890abcdef-01";
       vi.spyOn(InvokeStore, "getInstanceAsync").mockResolvedValue({
         getXRayTraceId: () => undefined,
@@ -277,8 +277,37 @@ describe(recursionDetectionMiddleware.name, () => {
 
       const { request } = mockNextHandler.mock.calls[0][0];
       expect(request.headers["traceparent"]).toBe(existingTraceparent);
-      expect(request.headers["tracestate"]).toBe(mockTracestate);
-      expect(request.headers["baggage"]).toBe(mockBaggage);
+      expect(request.headers["tracestate"]).toBeUndefined();
+      expect(request.headers["baggage"]).toBeUndefined();
+    });
+
+    it("should preserve existing traceparent, tracestate, and baggage headers without overwriting from InvokeStore", async () => {
+      const existingTraceparent = "00-abcdef1234567890abcdef1234567890-1234567890abcdef-01";
+      const existingTracestate = "vendor1=opaqueValue1";
+      const existingBaggage = "key1=value1";
+      vi.spyOn(InvokeStore, "getInstanceAsync").mockResolvedValue({
+        getXRayTraceId: () => undefined,
+        getTraceparent: () => mockTraceparent,
+        getTracestate: () => mockTracestate,
+        getBaggage: () => mockBaggage,
+      } as any);
+
+      const handler = recursionDetectionMiddleware()(mockNextHandler, {} as any);
+      await handler({
+        input: {},
+        request: new HttpRequest({
+          headers: {
+            traceparent: existingTraceparent,
+            tracestate: existingTracestate,
+            baggage: existingBaggage,
+          },
+        }),
+      });
+
+      const { request } = mockNextHandler.mock.calls[0][0];
+      expect(request.headers["traceparent"]).toBe(existingTraceparent);
+      expect(request.headers["tracestate"]).toBe(existingTracestate);
+      expect(request.headers["baggage"]).toBe(existingBaggage);
     });
 
     it("should not overwrite existing traceparent when InvokeStore also has one", async () => {

packages-internal/core/src/submodules/client/middleware-recursion-detection/recursionDetectionMiddleware.ts

@@ -65,19 +65,21 @@ export const recursionDetectionMiddleware =
 
       sanitizeTraceHeaders(request.headers);
 
-      const traceparent =
-        request.headers[TRACEPARENT] ?? (invokeStore ??= await InvokeStore.getInstanceAsync())?.getTraceparent();
-      if (traceparent) {
-        // traceparent is required.
-        request.headers[TRACEPARENT] = traceparent;
+      const existingTraceparent = request.headers[TRACEPARENT];
 
-        const tracestate = invokeStore?.getTracestate?.();
-        if (tracestate) {
-          request.headers[TRACESTATE] = tracestate;
-        }
-        const baggage = invokeStore?.getBaggage?.();
-        if (baggage) {
-          request.headers[BAGGAGE] = baggage;
+      if (!existingTraceparent) {
+        const traceparent = (invokeStore ??= await InvokeStore.getInstanceAsync())?.getTraceparent?.();
+        if (traceparent) {
+          request.headers[TRACEPARENT] = traceparent;
+
+          const tracestate = invokeStore?.getTracestate?.();
+          if (tracestate) {
+            request.headers[TRACESTATE] = tracestate;
+          }
+          const baggage = invokeStore?.getBaggage?.();
+          if (baggage) {
+            request.headers[BAGGAGE] = baggage;
+          }
         }
       }
     }