diff --git a/src/commands/stepfunctions/__tests__/helpers.test.ts b/src/commands/stepfunctions/__tests__/helpers.test.ts
index 8bcfc0ae3..598787cd5 100644
--- a/src/commands/stepfunctions/__tests__/helpers.test.ts
+++ b/src/commands/stepfunctions/__tests__/helpers.test.ts
@@ -53,7 +53,7 @@ describe('stepfunctions command helpers tests', () => {
         End: true,
       }
       expect(injectContextForLambdaFunctions(step, context, 'Lambda Invoke')).toBeTruthy()
-      expect(step.Parameters?.['Payload.$']).toEqual('States.JsonMerge($$, $, false)')
+      expect(step.Parameters?.['Payload.$']).toEqual(`$$['Execution', 'State']`)
     })
 
     test('default payload field of $', () => {
diff --git a/src/commands/stepfunctions/helpers.ts b/src/commands/stepfunctions/helpers.ts
index 0edd8238c..6a543a985 100644
--- a/src/commands/stepfunctions/helpers.ts
+++ b/src/commands/stepfunctions/helpers.ts
@@ -110,10 +110,6 @@ export const injectContextIntoTasks = async (
   }
 }
 
-export const addTraceContextToLambdaParameters = (Parameters: ParametersType): void => {
-  Parameters[`Payload.$`] = 'States.JsonMerge($$, $, false)'
-}
-
 export const addTraceContextToStepFunctionParameters = ({Parameters}: StepType): void => {
   if (Parameters) {
     if (!Parameters.Input) {
@@ -236,14 +232,14 @@ check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\
 
   // payload field not set
   if (!step.Parameters.hasOwnProperty('Payload.$') && !step.Parameters.hasOwnProperty('Payload')) {
-    addTraceContextToLambdaParameters(step.Parameters)
+    step.Parameters[`Payload.$`] = `$$['Execution', 'State']`
 
     return true
   }
 
   // default payload
   if (step.Parameters['Payload.$'] === '$') {
-    addTraceContextToLambdaParameters(step.Parameters)
+    step.Parameters[`Payload.$`] = 'States.JsonMerge($$, $, false)'
 
     return true
   }