[SVLS-518] fix: fix broken SFN -> lambda link (#779)

datadog-lambda-js ignores upstream trace context propagated via Execution.Input._datadog when a Lambda is invoked from a Step Function, breaking trace continuity for the pattern: upstream Lambda → StartExecution → Step Function → downstream Lambda. This commit extracts datadog headers from Execution.Input._datadog when it is present, linking a lambda span to its upstream step function.

Author: ida613

src/trace/context/extractors/step-function.spec.ts

@@ -45,6 +45,33 @@ describe("StepFunctionEventTraceExtractor", () => {
         Name: "abhinav-activity-state-machine",
       },
     };
+
+    const legacyLambdaRootPayload = {
+      Execution: {
+        Id: "arn:aws:states:sa-east-1:425362996713:execution:abhinav-activity-state-machine:72a7ca3e-901c-41bb-b5a3-5f279b92a316",
+        Name: "72a7ca3e-901c-41bb-b5a3-5f279b92a316",
+        RoleArn:
+          "arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j",
+        StartTime: "2024-12-04T19:38:04.069Z",
+        RedriveCount: 1,
+        Input: {
+          MyInput: "MyValue",
+          _datadog: {
+            "x-datadog-trace-id": "10593586103637578129",
+            "x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=6734e7c300000000",
+          },
+        },
+      },
+      State: {
+        Name: "Lambda Invoke",
+        EnteredTime: "2024-12-04T19:38:04.118Z",
+        RetryCount: 0,
+      },
+      StateMachine: {
+        Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-activity-state-machine",
+        Name: "abhinav-activity-state-machine",
+      },
+    };
     it("extracts trace context with valid payload", () => {
       // Mimick TraceContextService.extract initialization
       StepFunctionContextService.instance(payload);
@@ -61,7 +88,21 @@ describe("StepFunctionEventTraceExtractor", () => {
       expect(traceContext?.source).toBe("event");
     });
 
-    // https://github.com/DataDog/logs-backend/blob/c17618cb552fc369ca40282bae0a65803f82f694/domains/serverless/apps/logs-to-traces-reducer/src/test/resources/test-json-files/stepfunctions/RedriveTest/snapshots/RedriveLambdaSuccessTraceMerging.json#L46
+    it("extracts trace context with valid legacy lambda root payload", () => {
+      // Mimick TraceContextService.extract initialization
+      StepFunctionContextService.instance(legacyLambdaRootPayload);
+
+      const extractor = new StepFunctionEventTraceExtractor();
+
+      // Payload is sent again for safety in case the instance wasn't previously initialized
+      const traceContext = extractor.extract(legacyLambdaRootPayload);
+      expect(traceContext).not.toBeNull();
+
+      expect(traceContext?.toTraceId()).toBe("10593586103637578129");
+      expect(traceContext?.sampleMode()).toBe("1");
+      expect(traceContext?.source).toBe("event");
+    });
+
     it("extracts trace context with valid redriven payload", () => {
       // Mimick TraceContextService.extract initialization
       StepFunctionContextService.instance(redrivePayload);

src/trace/step-function-service.spec.ts

@@ -22,6 +22,33 @@ describe("StepFunctionContextService", () => {
       Name: "my-state-machine",
     },
   } as const;
+  const legacyLambdaRootStepFunctionEvent = {
+    _datadog: {
+      Execution: {
+        Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
+        Input: {
+          MyInput: "MyValue",
+          _datadog: {
+            "x-datadog-trace-id": "10593586103637578129",
+            "x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=6734e7c300000000",
+          },
+        },
+        Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be",
+        RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
+        RedriveCount: 0,
+        StartTime: "2022-12-08T21:08:17.924Z",
+      },
+      State: {
+        Name: "step-one",
+        EnteredTime: "2022-12-08T21:08:19.224Z",
+        RetryCount: 2,
+      },
+      StateMachine: {
+        Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
+        Name: "my-state-machine",
+      },
+    },
+  } as const;
   const lambdaRootStepFunctionEvent = {
     _datadog: {
       Execution: {
@@ -156,6 +183,23 @@ describe("StepFunctionContextService", () => {
       });
     });
 
+    it("sets context from valid legacy lambda root step function event", () => {
+      const instance = StepFunctionContextService.instance();
+      // Force setting event
+      instance["setContext"](legacyLambdaRootStepFunctionEvent);
+      expect(instance.context).toEqual({
+        execution_id:
+          "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
+        redrive_count: "0",
+        retry_count: "2",
+        state_entered_time: "2022-12-08T21:08:19.224Z",
+        state_name: "step-one",
+        trace_id: "10593586103637578129",
+        dd_p_tid: "6734e7c300000000",
+        serverless_version: "legacy-lambda-root",
+      });
+    });
+
     it("sets context from valid nested event", () => {
       const instance = StepFunctionContextService.instance();
       // Force setting event
@@ -211,6 +255,21 @@ describe("StepFunctionContextService", () => {
       expect(spanContext?.source).toBe("event");
     });
 
+    it("returns a SpanContextWrapper when legacy lambda root step function event is valid", () => {
+      const instance = StepFunctionContextService.instance();
+      // Force setting event
+      instance["setContext"](legacyLambdaRootStepFunctionEvent);
+
+      const spanContext = instance.spanContext;
+
+      expect(spanContext).not.toBeNull();
+
+      expect(spanContext?.toTraceId()).toBe("10593586103637578129");
+      expect(spanContext?.toSpanId()).toBe("7747304477664363642");
+      expect(spanContext?.sampleMode()).toBe("1");
+      expect(spanContext?.source).toBe("event");
+    });
+
     it("returns a SpanContextWrapper when nested event is valid", () => {
       const instance = StepFunctionContextService.instance();
       // Force setting event

src/trace/step-function-service.ts

@@ -120,13 +120,32 @@ export class StepFunctionContextService {
         } as LambdaRootStepFunctionContext;
       }
     } else {
-      this.context = {
-        execution_id,
-        redrive_count,
-        retry_count,
-        state_entered_time,
-        state_name,
-      } as LegacyStepFunctionContext;
+      const datadogContext = event.Execution?.Input?._datadog;
+      if (
+        typeof datadogContext === "object" &&
+        datadogContext !== null &&
+        typeof datadogContext["x-datadog-trace-id"] === "string" &&
+        typeof datadogContext["x-datadog-tags"] === "string"
+      ) {
+        this.context = {
+          execution_id,
+          redrive_count,
+          retry_count,
+          state_entered_time,
+          state_name,
+          trace_id: datadogContext["x-datadog-trace-id"],
+          dd_p_tid: this.parsePTid(datadogContext["x-datadog-tags"]),
+          serverless_version: "legacy-lambda-root", // dummy value
+        } as LambdaRootStepFunctionContext;
+      } else {
+        this.context = {
+          execution_id,
+          redrive_count,
+          retry_count,
+          state_entered_time,
+          state_name,
+        } as LegacyStepFunctionContext;
+      }
     }
   }