feat(stepfunctions-tasks): support Associate Workflow Executions on StepFunctionsStartExecution via associateWithParent property

packages/@aws-cdk/aws-stepfunctions-tasks/README.md

@@ -1111,6 +1111,21 @@ new sfn.StateMachine(this, 'ParentStateMachine', {
 });
 ```
 
+You can utilize [Associate Workflow Executions](https://docs.aws.amazon.com/step-functions/latest/dg/concepts-nested-workflows.html#nested-execution-startid)
+via the `associateWithParent` property. This allows the Step Functions UI to link child
+executions from parent executions, making it easier to trace execution flow across state machines.
+
+```ts
+const task = new tasks.StepFunctionsStartExecution(this, 'ChildTask', {
+  stateMachine: child,
+  associateWithParent: true,
+});
+```
+
+This will add the payload `AWS_STEP_FUNCTIONS_STARTED_BY_EXECUTION_ID.$: $$.Execution.Id` to the
+`input`property for you, which will pass the execution ID from the context object to the
+execution input. It requires `input` to be an object or not be set at all.
+
 ### Invoke Activity
 
 You can invoke a [Step Functions Activity](https://docs.aws.amazon.com/step-functions/latest/dg/concepts-activities.html) which enables you to have

packages/@aws-cdk/aws-stepfunctions-tasks/lib/stepfunctions/start-execution.ts

@@ -30,6 +30,18 @@ export interface StepFunctionsStartExecutionProps extends sfn.TaskStateBaseProps
    * @default - None
    */
   readonly name?: string;
+
+  /**
+   * Pass the execution ID from the context object to the execution input.
+   * This allows the Step Functions UI to link child executions from parent executions, making it easier to trace execution flow across state machines.
+   *
+   * If you set this property to `true`, the `input` property must be an object (provided by `sfn.TaskInput.fromObject`) or omitted entirely.
+   *
+   * @see https://docs.aws.amazon.com/step-functions/latest/dg/concepts-nested-workflows.html#nested-execution-startid
+   *
+   * @default - false
+   */
+  readonly associateWithParent?: boolean;
 }
 
 /**
@@ -59,6 +71,10 @@ export class StepFunctionsStartExecution extends sfn.TaskStateBase {
       throw new Error('Task Token is required in `input` for callback. Use JsonPath.taskToken to set the token.');
     }
 
+    if (this.props.associateWithParent && props.input && props.input.type !== sfn.InputType.OBJECT) {
+      throw new Error('Could not enable `associateWithParent` because `input` is taken directly from a JSON path. Use `sfn.TaskInput.fromObject` instead.');
+    }
+
     this.taskPolicies = this.createScopedAccessPolicy();
   }
 
@@ -70,10 +86,20 @@ export class StepFunctionsStartExecution extends sfn.TaskStateBase {
     // suffix is only applicable when waiting for a nested state machine to complete (RUN_JOB)
     // https://docs.aws.amazon.com/step-functions/latest/dg/connect-stepfunctions.html
     const suffix = this.integrationPattern === sfn.IntegrationPattern.RUN_JOB ? ':2' : '';
+    let input: any;
+    if (this.props.associateWithParent) {
+      const associateWithParentEntry = {
+        AWS_STEP_FUNCTIONS_STARTED_BY_EXECUTION_ID: sfn.JsonPath.stringAt('$$.Execution.Id'),
+      };
+      input = this.props.input ? { ...this.props.input.value, ...associateWithParentEntry } : associateWithParentEntry;
+    } else {
+      input = this.props.input ? this.props.input.value: sfn.TaskInput.fromJsonPathAt('$').value;
+    }
+
     return {
       Resource: `${integrationResourceArn('states', 'startExecution', this.integrationPattern)}${suffix}`,
       Parameters: sfn.FieldUtils.renderObject({
-        Input: this.props.input ? this.props.input.value : sfn.TaskInput.fromJsonPathAt('$').value,
+        Input: input,
         StateMachineArn: this.props.stateMachine.stateMachineArn,
         Name: this.props.name,
       }),

packages/@aws-cdk/aws-stepfunctions-tasks/test/stepfunctions/start-execution.test.ts

@@ -215,3 +215,55 @@ test('Execute State Machine - Wait For Task Token - Missing Task Token', () => {
     });
   }).toThrow('Task Token is required in `input` for callback. Use JsonPath.taskToken to set the token.');
 });
+
+test('Execute State Machine - Associate With Parent - Input Provided', () => {
+  const task = new StepFunctionsStartExecution(stack, 'ChildTask', {
+    stateMachine: child,
+    input: sfn.TaskInput.fromObject({
+      token: sfn.JsonPath.taskToken,
+    }),
+    associateWithParent: true,
+  });
+
+  new sfn.StateMachine(stack, 'ParentStateMachine', {
+    definition: task,
+  });
+
+  expect(stack.resolve(task.toStateJson())).toMatchObject({
+    Parameters: {
+      Input: {
+        'token.$': '$$.Task.Token',
+        'AWS_STEP_FUNCTIONS_STARTED_BY_EXECUTION_ID.$': '$$.Execution.Id',
+      },
+    },
+  });
+});
+
+test('Execute State Machine - Associate With Parent - Input Not Provided', () => {
+  const task = new StepFunctionsStartExecution(stack, 'ChildTask', {
+    stateMachine: child,
+    associateWithParent: true,
+  });
+
+  new sfn.StateMachine(stack, 'ParentStateMachine', {
+    definition: task,
+  });
+
+  expect(stack.resolve(task.toStateJson())).toMatchObject({
+    Parameters: {
+      Input: {
+        'AWS_STEP_FUNCTIONS_STARTED_BY_EXECUTION_ID.$': '$$.Execution.Id',
+      },
+    },
+  });
+});
+
+test('Execute State Machine - Associate With Parent - Incorrect Input Type', () => {
+  expect(() => {
+    new StepFunctionsStartExecution(stack, 'ChildTask', {
+      stateMachine: child,
+      associateWithParent: true,
+      input: sfn.TaskInput.fromText('{ "token.$": "$$.Task.Token" }'),
+    });
+  }).toThrow('Could not enable `associateWithParent` because `input` is taken directly from a JSON path. Use `sfn.TaskInput.fromObject` instead.');
+});