From 95a8379854e2bda6db442af81c629d65929290a4 Mon Sep 17 00:00:00 2001
From: Yiming Luo <yiming.luo@datadoghq.com>
Date: Thu, 26 Sep 2024 11:41:19 -0400
Subject: [PATCH] feat: Support sfn->Lambda context injection when Payload.$
 exists (case 4)

---
 src/step-functions-helper.spec.ts | 34 +++++++++++++++++++++++++++--
 src/step-functions-helper.ts      | 36 +++++++++++++++++++++++--------
 2 files changed, 59 insertions(+), 11 deletions(-)

diff --git a/src/step-functions-helper.spec.ts b/src/step-functions-helper.spec.ts
index 7efc1b3..a5d2e53 100644
--- a/src/step-functions-helper.spec.ts
+++ b/src/step-functions-helper.spec.ts
@@ -51,7 +51,7 @@ describe("test updateDefinitionString", () => {
     expect(definitionAfterUpdate.States?.InvokeLambda?.Parameters).toBe("Just a string!");
   });
 
-  it("test lambda step with default payload of '$'", async () => {
+  it("Case 4.1: test lambda step with default payload of '$'", async () => {
     const definitionString = {
       "Fn::Sub": [
         '{"Comment":"fake comment","StartAt":"InvokeLambda","States":{"InvokeLambda":{"Type":"Task","Parameters":{"FunctionName":"fake-function-name","Payload.$":"$"},"Resource":"arn:aws:states:::lambda:invoke","End":true}}}',
@@ -179,7 +179,37 @@ describe("test updateDefinitionString", () => {
     });
   });
 
-  it("test lambda step already has customized payload set do nothing", async () => {
+  it(`Case 4.2: test lambda step already has context injection set up using "Payload.$": "States.JsonMerge($$, $, false)"`, async () => {
+    const definitionString = {
+      "Fn::Sub": [
+        '{"Comment":"fake comment","StartAt":"InvokeLambda","States":{"InvokeLambda":{"Type":"Task","Parameters":{"FunctionName":"fake-function-name","Payload.$":"States.JsonMerge($$, $, false)"},"Resource":"arn:aws:states:::lambda:invoke","End":true}}}',
+        {},
+      ],
+    };
+    updateDefinitionString(definitionString, serverless, stateMachineName);
+
+    const definitionAfterUpdate: StateMachineDefinition = JSON.parse(definitionString["Fn::Sub"][0] as string);
+    expect(definitionAfterUpdate.States?.InvokeLambda?.Parameters?.["Payload.$"]).toBe(
+      "States.JsonMerge($$, $, false)",
+    );
+  });
+
+  it(`Case 4.2: test lambda step already has context injection set up using "Payload.$": "$$['Execution', 'State', 'StateMachine']"`, async () => {
+    const definitionString = {
+      "Fn::Sub": [
+        `{"Comment":"fake comment","StartAt":"InvokeLambda","States":{"InvokeLambda":{"Type":"Task","Parameters":{"FunctionName":"fake-function-name","Payload.$":"$$['Execution', 'State', 'StateMachine']"},"Resource":"arn:aws:states:::lambda:invoke","End":true}}}`,
+        {},
+      ],
+    };
+    updateDefinitionString(definitionString, serverless, stateMachineName);
+
+    const definitionAfterUpdate: StateMachineDefinition = JSON.parse(definitionString["Fn::Sub"][0] as string);
+    expect(definitionAfterUpdate.States?.InvokeLambda?.Parameters?.["Payload.$"]).toBe(
+      `$$['Execution', 'State', 'StateMachine']`,
+    );
+  });
+
+  it("Case 4.3: test lambda step has custom Payload.$ do nothing", async () => {
     const definitionString = {
       "Fn::Sub": [
         '{"Comment":"fake comment","StartAt":"InvokeLambda","States":{"InvokeLambda":{"Type":"Task","Parameters":{"FunctionName":"fake-function-name","Payload.$":"something-customized"},"Resource":"arn:aws:states:::lambda:invoke","End":true}}}',
diff --git a/src/step-functions-helper.ts b/src/step-functions-helper.ts
index 918040b..7ed259a 100644
--- a/src/step-functions-helper.ts
+++ b/src/step-functions-helper.ts
@@ -207,20 +207,38 @@ merge these traces, check out https://docs.datadoghq.com/serverless/step_functio
 
       return;
     }
-  }
+  } else {
+    // Case 4: Parameters has "Payload.$" field
+
+    // Case 4.1: default "Payload.$"
+    if (step.Parameters["Payload.$"] === "$") {
+      step.Parameters!["Payload.$"] = "States.JsonMerge($$, $, false)";
+      serverless.cli.log(
+        `JsonMerge Step Functions context object with payload in step: ${stepName} of state machine: ${stateMachineName}.`,
+      );
+      return;
+    }
+
+    // Case 4.2: context injection is already set up using "Payload.$"
+    if (
+      step.Parameters["Payload.$"] === "States.JsonMerge($$, $, false)" ||
+      step.Parameters["Payload.$"] === "$$['Execution', 'State', 'StateMachine']"
+    ) {
+      serverless.cli.log(
+        `[Warn] Step ${stepName} of state machine ${stateMachineName}: Context injection is already set up. Skipping context injection.\n`,
+      );
 
-  if (step.Parameters["Payload.$"] === "$") {
-    // $ % is the default original unchanged payload
-    step.Parameters!["Payload.$"] = "States.JsonMerge($$, $, false)";
+      return;
+    }
+
+    // Case 4.3: custom "Payload.$"
     serverless.cli.log(
-      `JsonMerge Step Functions context object with payload in step: ${stepName} of state machine: ${stateMachineName}.`,
+      `[Warn] Step ${stepName} of state machine ${stateMachineName} has a custom Payload field. Step Functions Context Object injection \
+skipped. Your Step Functions trace will not be merged with downstream Lambda traces. To manually merge these traces, \
+check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\n`,
     );
     return;
   }
-
-  serverless.cli.log(
-    `[Warn] Parameters.Payload has been set. Merging traces failed for step: ${stepName} of state machine: ${stateMachineName}`,
-  );
 }
 
 function updateDefinitionForStepFunctionInvocationStep(step: StateMachineStep): void {