diff --git a/src/step-functions-helper.spec.ts b/src/step-functions-helper.spec.ts
index ca2bd900..3dd3570a 100644
--- a/src/step-functions-helper.spec.ts
+++ b/src/step-functions-helper.spec.ts
@@ -1,5 +1,5 @@
 import {
-  isDefaultLambdaApiStep,
+  isLambdaApiStep,
   StateMachineDefinition,
   updateDefinitionString,
   updateDefinitionForStepFunctionInvocationStep,
@@ -223,7 +223,7 @@ describe("test updateDefinitionString", () => {
     expect(definitionAfterUpdate.States?.InvokeLambda?.Parameters?.["Payload.$"]).toBe("something-customized");
   });
 
-  it("test lambda basic legacy integration do nothing", async () => {
+  it("test lambda legacy integration with undefined parameters do nothing", async () => {
     const definitionString = {
       "Fn::Sub": [
         '{"Comment":"fake comment","StartAt":"InvokeLambda","States":{"InvokeLambda":{"Type":"Task","Resource":"arn:aws:lambda:sa-east-1:601427271234:function:unit-test-function-name","End":true}}}',
@@ -254,7 +254,7 @@ describe("test updateDefinitionString", () => {
     });
   });
 
-  it("test legacy lambda api do nothing", async () => {
+  it("test legacy lambda context is injected", async () => {
     const definitionString = {
       "Fn::Sub": [
         '{"Comment":"fake comment","StartAt":"InvokeLambda","States":{"InvokeLambda":{"Type":"Task","Parameters":{"FunctionName":"fake-function-name","Payload.$":"$"},"Resource":"arn:aws:lambda:sa-east-1:601427271234:function:unit-test-function-name","End":true}}}',
@@ -268,7 +268,7 @@ describe("test updateDefinitionString", () => {
       End: true,
       Parameters: {
         FunctionName: "fake-function-name",
-        "Payload.$": "$",
+        "Payload.$": "States.JsonMerge($$, $, false)",
       },
       Resource: "arn:aws:lambda:sa-east-1:601427271234:function:unit-test-function-name",
       Type: "Task",
@@ -381,29 +381,29 @@ describe("test updateDefinitionForStepFunctionInvocationStep", () => {
   });
 });
 
-describe("test isDefaultLambdaApiStep", () => {
+describe("test isLambdaApiStep", () => {
   it("resource is default lambda", async () => {
     const resource = "arn:aws:states:::lambda:invoke";
-    expect(isDefaultLambdaApiStep(resource)).toBeTruthy();
+    expect(isLambdaApiStep(resource)).toBeTruthy();
   });
 
   it("resource is lambda arn for legacy lambda api", async () => {
     const resource = "arn:aws:lambda:sa-east-1:601427271234:function:hello-function";
-    expect(isDefaultLambdaApiStep(resource)).toBeFalsy();
+    expect(isLambdaApiStep(resource)).toBeTruthy();
   });
 
   it("resource of dynamodb", async () => {
     const resource = "arn:aws:states:::dynamodb:updateItem";
-    expect(isDefaultLambdaApiStep(resource)).toBeFalsy();
+    expect(isLambdaApiStep(resource)).toBeFalsy();
   });
 
   it("resource of empty string", async () => {
     const resource = "";
-    expect(isDefaultLambdaApiStep(resource)).toBeFalsy();
+    expect(isLambdaApiStep(resource)).toBeFalsy();
   });
 
   it("resource of undefined", async () => {
     const resource = undefined;
-    expect(isDefaultLambdaApiStep(resource)).toBeFalsy();
+    expect(isLambdaApiStep(resource)).toBeFalsy();
   });
 });
diff --git a/src/step-functions-helper.ts b/src/step-functions-helper.ts
index ec5eab89..77025607 100644
--- a/src/step-functions-helper.ts
+++ b/src/step-functions-helper.ts
@@ -42,16 +42,12 @@ export interface StateMachineStep {
   End?: boolean;
 }
 
-export function isDefaultLambdaApiStep(resource: string | undefined): boolean {
-  // default means not legacy lambda api
-  if (resource === undefined) {
-    return false;
-  }
-  if (resource === "arn:aws:states:::lambda:invoke") {
-    // Legacy Lambda API resource.startsWith("arn:aws:lambda"), but it cannot inject context obj.
-    return true;
-  }
-  return false;
+export function isLambdaApiStep(resource: string | undefined): boolean {
+  // Allow for either the standard or legacy definitions of a lambda step
+  return (
+    resource !== undefined &&
+    (resource?.startsWith("arn:aws:states:::lambda:invoke") || resource?.startsWith("arn:aws:lambda"))
+  );
 }
 
 export function isStepFunctionInvocation(resource: string | undefined): boolean {
@@ -99,8 +95,8 @@ export function updateDefinitionString(
   // Step 2: Mutate the definition object
   const states = definitionObj.States;
   for (const [stepName, step] of Object.entries(states)) {
-    // only inject context into default Lambda API steps and Step Function invocation steps
-    if (isDefaultLambdaApiStep(step?.Resource)) {
+    // only inject context into Lambda API steps and Step Function invocation steps
+    if (isLambdaApiStep(step?.Resource)) {
       updateDefinitionForDefaultLambdaApiStep(stepName, step, serverless, stateMachineName);
     } else if (isStepFunctionInvocation(step?.Resource)) {
       updateDefinitionForStepFunctionInvocationStep(stepName, step, serverless, stateMachineName);