AWS Step Functions CDK IAM role circular dependency with distributed map state

I'm using AWS Step Functions with CDK (TypeScript) and I need to add a Distributed Map state. Unfortunately, CDK does not support this state yet (https://github.com/aws/aws-cdk/issues/23216).

Since I would like to create something a bit more structured than the CustomState proposed in the issue, I've made this class:

import { Construct } from "constructs/lib";
import { Map, MapProps } from 'aws-cdk-lib/aws-stepfunctions';

export enum ExecutionType {
  STANDARD = "STANDARD",
  EXPRESS = "EXPRESS"
}

export interface DistributedMapProps extends MapProps {
  readonly executionType?: ExecutionType
}

//https://github.com/aws/aws-cdk/issues/23216
export class DistributedMap extends Map {

  private executionType: ExecutionType
  private distributedMaxConcurrency?: number

  constructor(scope: Construct, id: string, props: DistributedMapProps = {}) { 
    super(scope, id, props)
    this.executionType = props.executionType ?? ExecutionType.STANDARD
    this.distributedMaxConcurrency = props.maxConcurrency ?? 1000
  }

  public override toStateJson(): object {
    let json = super.toStateJson() as any

    json.Iterator = {
      ...json.Iterator,
      ProcessorConfig: {
        Mode: "DISTRIBUTED",
        ExecutionType: this.executionType
      }
    }

    json.MaxConcurrency = this.distributedMaxConcurrency
    return json
  }
}

As you can see it extends Map and it works exactly like a normal Map. It just add the ProcessorConfig node into the Iterator node when toStateJson() is called.

S3 Buckets operations are not supported by this class because I don't need them at the moment: I just need a normal Map loop with 1k parallel executions.

Well this works like a charm... Except for permissions on the IAM Role.

To execute a distributed MapState I have to grant the StartExecution to my step function of the step function itself ( https://docs.aws.amazon.com/step-functions/latest/dg/iam-policies-eg-dist-map.html)

Here's what I've tried:

//No error or warnings, but also no changes on the IAM Role
myStateMachine.grantStartExecution(myStateMachine)


//Circular reference error at Synth time
myStateMachine.addToRolePolicy(new PolicyStatement({
    effect: Effect.ALLOW,
    actions: ["states:StartExecution"],
    resources: [myStateMachine.stateMachineArn],
}))

I've also tried to create an extra Stack only to update the Role (after the StateMachine Stack execution), but I can't use addToRolePolicy since StateMachine.fromStateMachineName() returns an IStateMachine instead of the real StateMachine. Calling grantStartExecution on the IStateMachine produces a warning at Synth time and no changes on the role after the deploy.

In the GitHub issue this guy used this workaround (https://github.com/aws/aws-cdk/issues/23216#issuecomment-1633630213):

private fixDistributedMapIamPermission() {
    const dummyMapGraph = new sfn.StateGraph(this.dummyMap, "dummyMapGraph");
    dummyMapGraph.policyStatements.map(p => this.stepFunction.addToRolePolicy(p));
    
    // https://docs.aws.amazon.com/step-functions/latest/dg/iam-policies-eg-dist-map.html#iam-policy-run-dist-map
    this.stepFunction.addToRolePolicy(new iam.PolicyStatement({
      effect: iam.Effect.ALLOW,
      actions: ["states:StartExecution", "states:DescribeExecution", "states:StopExecution"],
      resources: ["*"],
    }));
  }

Since I'm not using the dummy map (I use hierarchy instead) I can't copy/paste this. I tried using directly my DistributedMap instead of the dummyMap (since it's a Map itself) but it doesn't compile because I can't create a StateGraph from a state already in use in my StateMachine. Anyway I don't really understand how this workaround should work.

Any idea about how to solve this? I would like to keep my DistributedMapState class mostly unchanged and find a way to update just the IAM Role (avoiding the circularity), but I can make changes freely if I have no other choices anyway.

Thank you in advance

I found the following solution, it worked without changing anything in my code except for the IAM Role policy creation

//Create a new policy
const policy = new Policy(this, `${options.name} self-execution policy`, {
    statements: [
        new PolicyStatement({
            actions: ['states:StartExecution'],
            resources: [ stateMachine.stateMachineArn ]
        })
    ]
})

//Attach the new policy to the state machine
policy.attachToRole(stateMachine.role)

In other words it seems that creating a new policy with the needed permission does not lead to a circular dependency.