cross invocation tracecontext propagation

src/trace/context/extractor.ts

@@ -5,9 +5,11 @@ import { XrayService } from "../xray-service";
 import {
   AppSyncEventTraceExtractor,
   CustomTraceExtractor,
+  DurableExecutionEventTraceExtractor,
   EventBridgeEventTraceExtractor,
   EventBridgeSQSEventTraceExtractor,
   HTTPEventTraceExtractor,
+  isDurableExecutionEvent,
   KinesisEventTraceExtractor,
   LambdaContextTraceExtractor,
   SNSEventTraceExtractor,
@@ -56,6 +58,9 @@ export class TraceContextExtractor {
       }
     }
 
+    // No stripping needed — trace context is stored in dedicated
+    // `_datadog_{N}` checkpoint operations.
+
     if (spanContext === null) {
       this.stepFunctionContextService = StepFunctionContextService.instance(event);
       if (this.stepFunctionContextService?.context) {
@@ -81,6 +86,9 @@ export class TraceContextExtractor {
   private getTraceEventExtractor(event: any): EventTraceExtractor | undefined {
     if (!event || typeof event !== "object") return;
 
+    // Check for durable execution event first (has DurableExecutionArn + CheckpointToken)
+    if (isDurableExecutionEvent(event)) return new DurableExecutionEventTraceExtractor(this.tracerWrapper);
+
     const headers = event.headers ?? event.multiValueHeaders;
     if (headers !== null && typeof headers === "object") {
       return new HTTPEventTraceExtractor(this.tracerWrapper, this.config.decodeAuthorizerContext);

src/trace/context/extractors/durable-execution.spec.ts

@@ -0,0 +1,148 @@
+import { createDurableExecutionRootSpan, DurableExecutionEventTraceExtractor } from "./durable-execution";
+import { TracerWrapper } from "../../tracer-wrapper";
+
+jest.mock("dd-trace", () => ({
+  startSpan: jest.fn(),
+}));
+
+function makeTracerWrapper(extractReturn: any = null): TracerWrapper {
+  return { extract: jest.fn().mockReturnValue(extractReturn) } as unknown as TracerWrapper;
+}
+
+describe("DurableExecutionEventTraceExtractor", () => {
+  const tracer = require("dd-trace");
+  const startSpanMock = tracer.startSpan as jest.Mock;
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it("delegates checkpoint headers to the standard propagator", () => {
+    const executionArn =
+      "arn:aws:lambda:us-east-2:123456789012:function:demo:$LATEST/durable-execution/demo/abc";
+
+    const checkpointHeaders = {
+      "x-datadog-trace-id": "149750110124521191",
+      "x-datadog-parent-id": "987654321012345678",
+      "x-datadog-sampling-priority": "1",
+    };
+
+    const event = {
+      DurableExecutionArn: executionArn,
+      CheckpointToken: "t-1",
+      InitialExecutionState: {
+        Operations: [
+          {
+            Id: "op-1",
+            Name: "_datadog_0",
+            Status: "SUCCEEDED",
+            StepDetails: {
+              Result: JSON.stringify(checkpointHeaders),
+            },
+          },
+        ],
+      },
+    };
+
+    const sentinelContext = { sentinel: true };
+    const tracerWrapper = makeTracerWrapper(sentinelContext);
+    const extractor = new DurableExecutionEventTraceExtractor(tracerWrapper);
+    const context = extractor.extract(event);
+
+    expect(tracerWrapper.extract).toHaveBeenCalledWith(checkpointHeaders);
+    expect(context).toBe(sentinelContext);
+  });
+
+  it("returns null when no checkpoint or upstream context exists", () => {
+    const tracerWrapper = makeTracerWrapper();
+    const extractor = new DurableExecutionEventTraceExtractor(tracerWrapper);
+
+    const context = extractor.extract({
+      DurableExecutionArn: "arn:aws:lambda:us-east-2:123:function:demo",
+      CheckpointToken: "t-empty",
+      InitialExecutionState: { Operations: [] },
+    });
+
+    expect(context).toBeNull();
+    expect(tracerWrapper.extract).not.toHaveBeenCalled();
+  });
+
+  it("creates durable root span only for first invocation", () => {
+    const executionArn =
+      "arn:aws:lambda:us-east-2:123456789012:function:demo:$LATEST/durable-execution/demo/first";
+
+    const spanContext: any = {
+      _spanId: null,
+      _parentId: null,
+      toTraceId: () => "1111111111111111111",
+      toSpanId: () => "2222222222222222222",
+    };
+    const span = {
+      context: () => spanContext,
+      finish: jest.fn(),
+    };
+    startSpanMock.mockReturnValue(span);
+
+    const firstInvocationEvent = {
+      DurableExecutionArn: executionArn,
+      CheckpointToken: "t-first",
+      InitialExecutionState: {
+        Operations: [
+          {
+            Id: "op-1",
+            Name: "input",
+            Status: "RUNNING",
+            StartTimestamp: 1710000000000,
+            ExecutionDetails: {
+              InputPayload: JSON.stringify({ hello: "world" }),
+            },
+          },
+        ],
+      },
+    };
+
+    const root = createDurableExecutionRootSpan(firstInvocationEvent, null);
+
+    expect(root).not.toBeNull();
+    expect(startSpanMock).toHaveBeenCalledTimes(1);
+  });
+
+  it("skips durable root span creation on replay invocations", () => {
+    const executionArn =
+      "arn:aws:lambda:us-east-2:123456789012:function:demo:$LATEST/durable-execution/demo/replay";
+
+    const replayEvent = {
+      DurableExecutionArn: executionArn,
+      CheckpointToken: "t-replay",
+      InitialExecutionState: {
+        Operations: [
+            {
+              Id: "op-1",
+              Name: "_datadog_0",
+              Status: "SUCCEEDED",
+              StepDetails: {
+                Result: JSON.stringify({
+                "x-datadog-trace-id": "149750110124521191",
+                "x-datadog-parent-id": "538591322263933970",
+                "x-datadog-sampling-priority": "1",
+              }),
+            },
+          },
+          {
+            Id: "op-2",
+            Name: "callback_step_prepare",
+            Status: "SUCCEEDED",
+          },
+        ],
+      },
+    };
+
+    const tracerWrapper = makeTracerWrapper({ source: "Event" });
+    const extractor = new DurableExecutionEventTraceExtractor(tracerWrapper);
+    const extracted = extractor.extract(replayEvent);
+    const root = createDurableExecutionRootSpan(replayEvent, extracted);
+
+    expect(root).toBeNull();
+    expect(startSpanMock).not.toHaveBeenCalled();
+  });
+});