From 1149114a91f16ba1e3d9eb393f3c4c95bf15452c Mon Sep 17 00:00:00 2001 From: bgagent Date: Mon, 22 Jun 2026 23:14:24 -0500 Subject: [PATCH] refactor(cdk): collapse AgentSessionRole trust+grant into admitComputeRole (#213) Replace split addAssumingRole/grantAssumeToComputeRole with a single method so SessionRole admission always wires both IAM halves together. Co-authored-by: Cursor --- cdk/src/constructs/agent-session-role.ts | 41 ++++++------ cdk/src/constructs/ecs-agent-cluster.ts | 3 +- cdk/src/stacks/agent.ts | 1 - .../constructs/agent-session-role.test.ts | 67 +++++++++++++------ cdk/test/constructs/ecs-agent-cluster.test.ts | 7 +- 5 files changed, 71 insertions(+), 48 deletions(-) diff --git a/cdk/src/constructs/agent-session-role.ts b/cdk/src/constructs/agent-session-role.ts index d89dd527..95df0840 100644 --- a/cdk/src/constructs/agent-session-role.ts +++ b/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'], diff --git a/cdk/src/constructs/ecs-agent-cluster.ts b/cdk/src/constructs/ecs-agent-cluster.ts index fac2102a..1ccbf7c4 100644 --- a/cdk/src/constructs/ecs-agent-cluster.ts +++ b/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); diff --git a/cdk/src/stacks/agent.ts b/cdk/src/stacks/agent.ts index 4ed6aac0..192496d3 100644 --- a/cdk/src/stacks/agent.ts +++ b/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 diff --git a/cdk/test/constructs/agent-session-role.test.ts b/cdk/test/constructs/agent-session-role.test.ts index 103fe5e5..1f990a63 100644 --- a/cdk/test/constructs/agent-session-role.test.ts +++ b/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(); }); }); diff --git a/cdk/test/constructs/ecs-agent-cluster.test.ts b/cdk/test/constructs/ecs-agent-cluster.test.ts index c7f68711..a6164a56 100644 --- a/cdk/test/constructs/ecs-agent-cluster.test.ts +++ b/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