From e191be2d80c3e57c29c374aa41b6eb4948970e16 Mon Sep 17 00:00:00 2001
From: Yiming Luo <yiming.luo@datadoghq.com>
Date: Mon, 16 Sep 2024 12:16:30 -0400
Subject: [PATCH] fix: do not print warning if context injection is already set
 up

---
 src/commands/stepfunctions/helpers.ts | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/src/commands/stepfunctions/helpers.ts b/src/commands/stepfunctions/helpers.ts
index a4517e4e3..d17c05283 100644
--- a/src/commands/stepfunctions/helpers.ts
+++ b/src/commands/stepfunctions/helpers.ts
@@ -150,6 +150,13 @@ check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\
     return true
   }
 
+  // context injection is already set up
+  if (step.Parameters['Payload.$'] === 'States.JsonMerge($$, $, false)') {
+    context.stdout.write(` Step ${stepName}: Context injection is already set up. Skipping context injection.\n`)
+
+    return false
+  }
+
   // custom payload
   context.stdout
     .write(`[Warn] Step ${stepName} has a custom Payload field. Step Functions Context Object injection skipped. \
@@ -203,6 +210,13 @@ merge these traces, check out https://docs.datadoghq.com/serverless/step_functio
     return true
   }
 
+  // context injection is already set up
+  if (step.Parameters.Input['CONTEXT.$'] === 'States.JsonMerge($$, $, false)') {
+    context.stdout.write(` Step ${stepName}: Context injection is already set up. Skipping context injection.\n`)
+
+    return false
+  }
+
   context.stdout
     .write(`[Warn] Step ${stepName}'s 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 \