refactor(cdk): collapse AgentSessionRole trust+grant into admitComputeRole (#213) (#430)

Replace split addAssumingRole/grantAssumeToComputeRole with a single method
so SessionRole admission always wires both IAM halves together.

Co-authored-by: bgagent <bgagent@noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

Author: 3 people

cdk/src/constructs/agent-session-role.ts

@@ -116,12 +116,12 @@ export class AgentSessionRole extends Construct {
       );
     }
 
-    const principals = props.assumingRoles.map(
-      (r) => new iam.ArnPrincipal(r.roleArn),
-    );
+    const [firstAssumingRole] = props.assumingRoles;
 
+    // CDK requires assumedBy; additional principals are admitted via
+    // admitComputeRole so trust + grant always wire together.
     this.role = new iam.Role(this, 'Role', {
-      assumedBy: new iam.CompositePrincipal(...principals),
+      assumedBy: new iam.ArnPrincipal(firstAssumingRole.roleArn),
       description:
         'Per-task scoped credentials for ABCA agent tenant-data access '
         + '(DynamoDB task rows + S3 trace/attachment objects), constrained by '
@@ -132,16 +132,6 @@ export class AgentSessionRole extends Construct {
       maxSessionDuration: Duration.hours(1),
     });
 
-    // The agent passes session tags on AssumeRole, which requires the trust
-    // policy to also allow sts:TagSession (CDK's assumedBy only adds
-    // sts:AssumeRole). Same principals.
-    this.role.assumeRolePolicy?.addStatements(
-      new iam.PolicyStatement({
-        actions: ['sts:TagSession'],
-        principals,
-      }),
-    );
-
     // --- DynamoDB: item access gated by task_id leading-key ---
     // One statement per table keeps the resource ARNs explicit. The condition
     // requires the request's partition key (task_id) to equal the session's
@@ -220,19 +210,28 @@ export class AgentSessionRole extends Construct {
       ],
       true,
     );
+
+    for (const computeRole of props.assumingRoles) {
+      this.admitComputeRole(computeRole);
+    }
   }
 
   /**
-   * Admit an additional compute role to the trust policy (e.g. the ECS Fargate
-   * task role when that backend is enabled). Adds `sts:AssumeRole` +
-   * `sts:TagSession` for the role to this SessionRole's trust policy.
+   * Admit a compute role to assume this SessionRole with session tags. Wires
+   * both halves AssumeRole requires: trust on this SessionRole and
+   * `sts:AssumeRole`/`sts:TagSession` on the compute role's identity policy.
    */
-  public addAssumingRole(computeRole: iam.IRole): void {
-    const principal = new iam.ArnPrincipal(computeRole.roleArn);
+  public admitComputeRole(computeRole: iam.IRole): void {
+    this.addTrustForComputeRole(computeRole);
+    this.grantAssumeToComputeRole(computeRole);
+  }
+
+  /** Add bundled AssumeRole + TagSession trust for one compute principal. */
+  private addTrustForComputeRole(computeRole: iam.IRole): void {
     this.role.assumeRolePolicy?.addStatements(
       new iam.PolicyStatement({
         actions: ['sts:AssumeRole', 'sts:TagSession'],
-        principals: [principal],
+        principals: [new iam.ArnPrincipal(computeRole.roleArn)],
       }),
     );
   }
@@ -243,7 +242,7 @@ export class AgentSessionRole extends Construct {
    * policy (a separate IAM::Policy resource, so no dependency cycle with this
    * role's trust policy).
    */
-  public grantAssumeToComputeRole(computeRole: iam.IRole): void {
+  private grantAssumeToComputeRole(computeRole: iam.IRole): void {
     computeRole.addToPrincipalPolicy(
       new iam.PolicyStatement({
         actions: ['sts:AssumeRole', 'sts:TagSession'],

cdk/src/constructs/ecs-agent-cluster.ts

@@ -147,8 +147,7 @@ export class EcsAgentCluster extends Construct {
     // that tag-scoped role and the task role only needs to assume it. Without
     // one (isolated construct tests / legacy), grant the task role directly.
     if (props.agentSessionRole) {
-      props.agentSessionRole.addAssumingRole(taskRole);
-      props.agentSessionRole.grantAssumeToComputeRole(taskRole);
+      props.agentSessionRole.admitComputeRole(taskRole);
     } else {
       props.taskTable.grantReadWriteData(taskRole);
       props.taskEventsTable.grantReadWriteData(taskRole);

cdk/src/stacks/agent.ts

@@ -478,7 +478,6 @@ export class AgentStack extends Stack {
       traceArtifactsBucket: traceArtifactsBucket.bucket,
       attachmentsBucket: attachmentsBucket.bucket,
     });
-    agentSessionRole.grantAssumeToComputeRole(runtime.role);
     sessionRoleArnHolder = agentSessionRole.role.roleArn;
 
     // X-Ray tracing disabled — requires account-level UpdateTraceSegmentDestination

cdk/test/constructs/agent-session-role.test.ts

@@ -59,11 +59,20 @@ function createStack() {
     traceArtifactsBucket,
     attachmentsBucket,
   });
-  sessionRole.grantAssumeToComputeRole(computeRole);
 
   return { stack, template: Template.fromStack(stack), computeRole, sessionRole };
 }
 
+function findComputeRolePolicy(template: Template) {
+  const policies = template.findResources('AWS::IAM::Policy');
+  return Object.entries(policies).find(([id]) => id.includes('ComputeRole'))![1];
+}
+
+function findSessionRoleTrust(template: Template) {
+  const roles = template.findResources('AWS::IAM::Role');
+  return Object.entries(roles).find(([id]) => id.includes('AgentSessionRole'))![1];
+}
+
 describe('AgentSessionRole construct', () => {
   let template: Template;
 
@@ -75,8 +84,9 @@ describe('AgentSessionRole construct', () => {
     template.hasResourceProperties('AWS::IAM::Role', {
       AssumeRolePolicyDocument: {
         Statement: Match.arrayWith([
-          Match.objectLike({ Action: 'sts:AssumeRole' }),
-          Match.objectLike({ Action: 'sts:TagSession' }),
+          Match.objectLike({
+            Action: Match.arrayWith(['sts:AssumeRole', 'sts:TagSession']),
+          }),
         ]),
       },
       // Role chaining caps at 1h regardless; documented explicitly.
@@ -163,20 +173,27 @@ describe('AgentSessionRole construct', () => {
     );
   });
 
-  test('compute role is granted sts:AssumeRole + sts:TagSession on the SessionRole', () => {
-    const policies = template.findResources('AWS::IAM::Policy');
-    const computePolicy = Object.entries(policies).find(([id]) =>
-      id.includes('ComputeRole'),
-    )![1];
+  test('constructor admits assumingRoles with both trust and grant wired', () => {
+    const computePolicy = findComputeRolePolicy(template);
     const statements = computePolicy.Properties.PolicyDocument.Statement;
     const stsStatement = statements.find((s: { Action: string | string[] }) => {
       const actions = Array.isArray(s.Action) ? s.Action : [s.Action];
       return actions.includes('sts:AssumeRole') && actions.includes('sts:TagSession');
     });
     expect(stsStatement).toBeDefined();
+
+    const trustDoc = findSessionRoleTrust(template).Properties.AssumeRolePolicyDocument;
+    const bundledTrust = trustDoc.Statement.find(
+      (s: { Action: string | string[] }) => {
+        const actions = Array.isArray(s.Action) ? s.Action : [s.Action];
+        return actions.includes('sts:AssumeRole') && actions.includes('sts:TagSession');
+      },
+    );
+    expect(bundledTrust).toBeDefined();
+    expect(JSON.stringify(trustDoc)).toContain('ComputeRole');
   });
 
-  test('addAssumingRole admits a second compute role (ECS task role) to the trust', () => {
+  test('admitComputeRole wires both trust and grant for an additional compute role', () => {
     const app = new App();
     const stack = new Stack(app, 'MultiPrincipalStack');
     const agentcoreRole = new iam.Role(stack, 'AgentCoreRole', {
@@ -194,17 +211,25 @@ describe('AgentSessionRole construct', () => {
       traceArtifactsBucket: new s3.Bucket(stack, 'TB'),
       attachmentsBucket: new s3.Bucket(stack, 'AB'),
     });
-    sessionRole.addAssumingRole(ecsTaskRole);
-
-    const trustStatements = Template.fromStack(stack).findResources(
-      'AWS::IAM::Role',
-    );
-    const sr = Object.entries(trustStatements).find(([id]) => id.includes('SR'))![1];
-    // Trust policy now has statements for both AssumeRole and TagSession; the
-    // ECS task role ARN appears as an additional principal.
-    const serialized = JSON.stringify(sr.Properties.AssumeRolePolicyDocument);
-    expect(serialized).toContain('sts:TagSession');
-    expect(serialized).toContain('EcsTaskRole');
-    expect(serialized).toContain('AgentCoreRole');
+    sessionRole.admitComputeRole(ecsTaskRole);
+
+    const stackTemplate = Template.fromStack(stack);
+    const trustDoc = Object.entries(stackTemplate.findResources('AWS::IAM::Role')).find(
+      ([id]) => id.includes('SR'),
+    )![1].Properties.AssumeRolePolicyDocument;
+    const serializedTrust = JSON.stringify(trustDoc);
+    expect(serializedTrust).toContain('sts:TagSession');
+    expect(serializedTrust).toContain('EcsTaskRole');
+    expect(serializedTrust).toContain('AgentCoreRole');
+
+    const ecsPolicy = Object.entries(stackTemplate.findResources('AWS::IAM::Policy')).find(
+      ([id]) => id.includes('EcsTaskRole'),
+    )![1];
+    const ecsStatements = ecsPolicy.Properties.PolicyDocument.Statement;
+    const stsGrant = ecsStatements.find((s: { Action: string | string[] }) => {
+      const actions = Array.isArray(s.Action) ? s.Action : [s.Action];
+      return actions.includes('sts:AssumeRole') && actions.includes('sts:TagSession');
+    });
+    expect(stsGrant).toBeDefined();
   });
 });

cdk/test/constructs/ecs-agent-cluster.test.ts

@@ -266,15 +266,16 @@ describe('EcsAgentCluster construct', () => {
       // to the SessionRole's policy (which carries the conditioned DDB
       // statements). The task-role policy must NOT contain any unconditioned
       // task-table DDB grant — that access now lives only on the SessionRole.
-      const taskRolePolicies = Object.values(policies).filter((p) =>
-        p.Properties.PolicyDocument.Statement.some((s: { Action: string | string[] }) => {
+      const taskRolePolicies = Object.entries(policies).filter(([id, p]) =>
+        id.includes('TaskDefTaskRole')
+        && p.Properties.PolicyDocument.Statement.some((s: { Action: string | string[] }) => {
           const actions = Array.isArray(s.Action) ? s.Action : [s.Action];
           return actions.includes('sts:AssumeRole');
         }),
       );
       expect(taskRolePolicies).toHaveLength(1);
 
-      const taskRoleStatements = taskRolePolicies[0].Properties.PolicyDocument.Statement;
+      const taskRoleStatements = taskRolePolicies[0][1].Properties.PolicyDocument.Statement;
       // No unconditioned dynamodb item grant on the task role (the only DDB the
       // task role may touch directly is UserConcurrencyTable — assert that any
       // DDB statement present is NOT a leading-key-less task-table grant by