From 07b5dd75445a1ed1ac0fc355ef88f41ff7186500 Mon Sep 17 00:00:00 2001
From: Yiming Luo <yiming.luo@datadoghq.com>
Date: Fri, 4 Oct 2024 16:07:19 -0400
Subject: [PATCH] Skip sfn->sfn context injection when CONTEXT.$ field exists
 (case 3) (#538)

---
 src/step-functions-helper.spec.ts | 14 +++++++++++++-
 src/step-functions-helper.ts      | 18 ++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/src/step-functions-helper.spec.ts b/src/step-functions-helper.spec.ts
index 5c287567..2af6d865 100644
--- a/src/step-functions-helper.spec.ts
+++ b/src/step-functions-helper.spec.ts
@@ -361,7 +361,19 @@ describe("test updateDefinitionForStepFunctionInvocationStep", () => {
     expect(updateDefinitionForStepFunctionInvocationStep(stepName, step, serverless, stateMachineName)).toBeFalsy();
   });
 
-  it("Input field has CONTEXT.$ already", async () => {
+  it("Case 3.1: Context injection already set up using States.JsonMerge($$, $, false)", async () => {
+    const parameters = { FunctionName: "bla", Input: { "CONTEXT.$": "States.JsonMerge($$, $, false)" } };
+    const step = { Parameters: parameters };
+    expect(updateDefinitionForStepFunctionInvocationStep(stepName, step, serverless, stateMachineName)).toBeFalsy();
+  });
+
+  it("Case 3.1: Context injection already set up using $$['Execution', 'State', 'StateMachine']", async () => {
+    const parameters = { FunctionName: "bla", Input: { "CONTEXT.$": "$$['Execution', 'State', 'StateMachine']" } };
+    const step = { Parameters: parameters };
+    expect(updateDefinitionForStepFunctionInvocationStep(stepName, step, serverless, stateMachineName)).toBeFalsy();
+  });
+
+  it("Case 3.2: Input field has a custom CONTEXT.$ field", async () => {
     const parameters = { FunctionName: "bla", Input: { "CONTEXT.$": "something else" } };
     const step = { Parameters: parameters };
     expect(updateDefinitionForStepFunctionInvocationStep(stepName, step, serverless, stateMachineName)).toBeFalsy();
diff --git a/src/step-functions-helper.ts b/src/step-functions-helper.ts
index 8a252118..cf138775 100644
--- a/src/step-functions-helper.ts
+++ b/src/step-functions-helper.ts
@@ -278,6 +278,24 @@ merge these traces, check out https://docs.datadoghq.com/serverless/step_functio
     return false;
   }
 
+  // Case 3.1 context injection is already set up
+  if (
+    parameters.Input["CONTEXT.$"] === "States.JsonMerge($$, $, false)" ||
+    parameters.Input["CONTEXT.$"] === `$$['Execution', 'State', 'StateMachine']`
+  ) {
+    serverless.cli.log(
+      `Step ${stepName} of state machine ${stateMachineName}: Context injection is already set up. Skipping context injection.\n`,
+    );
+
+    return false;
+  }
+
+  // Case 3.2 custom CONTEXT.$ field
+  serverless.cli
+    .log(`[Warn] Step ${stepName} of state machine ${stateMachineName}: Parameters.Input field has a custom CONTEXT.$ field. Step \
+Functions Context Object injection skipped. Your Step Functions trace will not be merged with downstream Step Function traces. To \
+manually merge these traces, check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\n`);
+
   return false;
 }