chore(pr): address Copilot review feedback

MichaelWalker-git · MichaelWalker-git · commit 42566d06b7ff · 2026-04-14T13:34:14.000-07:00
- Remove unnecessary iam:PassRole from orchestrator (EC2 strategy
  never passes a role to any API)
- Simplify ec2FleetConfig in task-api to empty object (instanceRoleArn
  was unused)
- Use CDK Tags.of() for ASG fleet tag propagation instead of no-op
  user-data tagging — instances are now tagged at launch
- Fix missing AWS_REGION in boot script by deriving from IMDS
- Eliminate shell injection risk by reading all task data from S3
  payload at runtime instead of interpolating into bash exports
- Add cleanup trap in boot script to always retag instance as idle
  on exit (success, error, or signal)
- Add try/catch rollback in startSession to retag instance as idle
  when SSM dispatch fails
- Generalize ECS-specific log messages in poll loop to be
  compute-backend-agnostic (uses strategy type label)
diff --git a/cdk/src/constructs/ec2-agent-fleet.ts b/cdk/src/constructs/ec2-agent-fleet.ts
@@ -17,7 +17,7 @@
  *  SOFTWARE.
  */
 
-import { Duration, RemovalPolicy } from 'aws-cdk-lib';
+import { Duration, RemovalPolicy, Tags } from 'aws-cdk-lib';
 import * as autoscaling from 'aws-cdk-lib/aws-autoscaling';
 import * as dynamodb from 'aws-cdk-lib/aws-dynamodb';
 import * as ec2 from 'aws-cdk-lib/aws-ec2';
@@ -184,10 +184,9 @@ export class Ec2AgentFleet extends Construct {
       healthCheck: autoscaling.HealthCheck.ec2(),
     });
 
-    // Tag the ASG instances for fleet identification
-    // CDK auto-propagates tags from the ASG to instances
-    this.autoScalingGroup.node.defaultChild;
-    this.autoScalingGroup.addUserData(`aws ec2 create-tags --resources "$(ec2-metadata -i | cut -d' ' -f2)" --region "$(ec2-metadata --availability-zone | cut -d' ' -f2 | sed 's/.$//')" --tags Key=${this.fleetTagKey},Value=${this.fleetTagValue}`);
+    // Tag ASG instances for fleet identification — CDK propagates these at launch
+    Tags.of(this.autoScalingGroup).add(this.fleetTagKey, this.fleetTagValue);
+    Tags.of(this.autoScalingGroup).add('bgagent:status', 'idle');
 
     NagSuppressions.addResourceSuppressions(this.instanceRole, [
       {
diff --git a/cdk/src/constructs/task-api.ts b/cdk/src/constructs/task-api.ts
@@ -111,9 +111,7 @@ export interface TaskApiProps {
    * EC2 fleet configuration for cancel-task to stop EC2-backed tasks.
    * When provided, the cancel Lambda gets `ssm:CancelCommand` permission.
    */
-  readonly ec2FleetConfig?: {
-    readonly instanceRoleArn: string;
-  };
+  readonly ec2FleetConfig?: Record<string, never>;
 }
 
 /**
diff --git a/cdk/src/constructs/task-orchestrator.ts b/cdk/src/constructs/task-orchestrator.ts
@@ -137,7 +137,6 @@ export interface TaskOrchestratorProps {
     readonly fleetTagValue: string;
     readonly payloadBucketName: string;
     readonly ecrImageUri: string;
-    readonly instanceRoleArn: string;
   };
 }
 
@@ -304,15 +303,9 @@ export class TaskOrchestrator extends Construct {
         resources: [`arn:${Aws.PARTITION}:s3:::${props.ec2Config.payloadBucketName}/*`],
       }));
 
-      this.fn.addToRolePolicy(new iam.PolicyStatement({
-        actions: ['iam:PassRole'],
-        resources: [props.ec2Config.instanceRoleArn],
-        conditions: {
-          StringEquals: {
-            'iam:PassedToService': 'ec2.amazonaws.com',
-          },
-        },
-      }));
+      // Note: iam:PassRole is not needed — the orchestrator does not pass the
+      // instance role to any EC2 API. The ASG launch template handles instance
+      // profile association at fleet creation time.
     }
 
     // Per-repo Secrets Manager grants (e.g. per-repo GitHub tokens from Blueprints)
diff --git a/cdk/src/handlers/orchestrate-task.ts b/cdk/src/handlers/orchestrate-task.ts
@@ -41,8 +41,8 @@ interface OrchestrateTaskEvent {
 
 const MAX_POLL_ATTEMPTS = 1020; // ~8.5h at 30s intervals
 const MAX_NON_RUNNING_POLLS = 10; // ~5min grace period for session to start
-const MAX_CONSECUTIVE_ECS_POLL_FAILURES = 3;
-const MAX_CONSECUTIVE_ECS_COMPLETED_POLLS = 5;
+const MAX_CONSECUTIVE_COMPUTE_POLL_FAILURES = 3;
+const MAX_CONSECUTIVE_COMPUTE_COMPLETED_POLLS = 5;
 
 const durableHandler: DurableExecutionHandler<OrchestrateTaskEvent, void> = async (event, context) => {
   const { task_id: taskId } = event;
@@ -176,63 +176,62 @@ const durableHandler: DurableExecutionHandler<OrchestrateTaskEvent, void> = asyn
     'await-agent-completion',
     async (state) => {
       const ddbState = await pollTaskStatus(taskId, state);
-      let consecutiveEcsPollFailures = 0;
-      let consecutiveEcsCompletedPolls = 0;
+      let consecutiveComputePollFailures = 0;
+      let consecutiveComputeCompletedPolls = 0;
+      const computeLabel = blueprintConfig.compute_type.toUpperCase();
 
-      // ECS compute-level crash detection: if DDB is not terminal, check ECS task status
+      // Compute-level crash detection: if DDB is not terminal, check compute task status
       if (
         ddbState.lastStatus &&
         !TERMINAL_STATUSES.includes(ddbState.lastStatus) &&
         computeStrategy
       ) {
         try {
-          const ecsStatus = await computeStrategy.pollSession(sessionHandle);
-          if (ecsStatus.status === 'failed') {
-            const errorMsg = 'error' in ecsStatus ? ecsStatus.error : 'ECS task failed';
-            logger.warn('ECS task failed before DDB terminal write', {
+          const computeStatus = await computeStrategy.pollSession(sessionHandle);
+          if (computeStatus.status === 'failed') {
+            const errorMsg = 'error' in computeStatus ? computeStatus.error : `${computeLabel} task failed`;
+            logger.warn(`${computeLabel} task failed before DDB terminal write`, {
               task_id: taskId,
               error: errorMsg,
             });
-            await failTask(taskId, ddbState.lastStatus, `ECS container failed: ${errorMsg}`, task.user_id, true);
+            await failTask(taskId, ddbState.lastStatus, `${computeLabel} compute failed: ${errorMsg}`, task.user_id, true);
             return { attempts: ddbState.attempts, lastStatus: TaskStatus.FAILED };
           }
-          if (ecsStatus.status === 'completed') {
-            consecutiveEcsCompletedPolls = (state.consecutiveEcsCompletedPolls ?? 0) + 1;
-            if (consecutiveEcsCompletedPolls >= MAX_CONSECUTIVE_ECS_COMPLETED_POLLS) {
-              // ECS task exited successfully but DDB never reached terminal — the agent
-              // likely crashed after container exit code 0 but before writing status.
-              logger.error('ECS task completed but DDB never caught up — failing task', {
+          if (computeStatus.status === 'completed') {
+            consecutiveComputeCompletedPolls = (state.consecutiveComputeCompletedPolls ?? 0) + 1;
+            if (consecutiveComputeCompletedPolls >= MAX_CONSECUTIVE_COMPUTE_COMPLETED_POLLS) {
+              logger.error(`${computeLabel} task completed but DDB never caught up — failing task`, {
                 task_id: taskId,
-                consecutive_completed_polls: consecutiveEcsCompletedPolls,
+                consecutive_completed_polls: consecutiveComputeCompletedPolls,
               });
-              await failTask(taskId, ddbState.lastStatus, `ECS task exited successfully but agent never wrote terminal status after ${consecutiveEcsCompletedPolls} polls`, task.user_id, true);
+              await failTask(taskId, ddbState.lastStatus, `${computeLabel} task exited successfully but agent never wrote terminal status after ${consecutiveComputeCompletedPolls} polls`, task.user_id, true);
               return { attempts: ddbState.attempts, lastStatus: TaskStatus.FAILED };
             }
-            logger.warn('ECS task completed but DDB not terminal — waiting for DDB catchup', {
+            logger.warn(`${computeLabel} task completed but DDB not terminal — waiting for DDB catchup`, {
               task_id: taskId,
-              consecutive_completed_polls: consecutiveEcsCompletedPolls,
+              consecutive_completed_polls: consecutiveComputeCompletedPolls,
             });
           }
         } catch (err) {
-          consecutiveEcsPollFailures = (state.consecutiveEcsPollFailures ?? 0) + 1;
-          if (consecutiveEcsPollFailures >= MAX_CONSECUTIVE_ECS_POLL_FAILURES) {
-            logger.error('ECS pollSession failed repeatedly — failing task', {
+          consecutiveComputePollFailures = (state.consecutiveComputePollFailures ?? 0) + 1;
+          if (consecutiveComputePollFailures >= MAX_CONSECUTIVE_COMPUTE_POLL_FAILURES) {
+            logger.error(`${computeLabel} pollSession failed repeatedly — failing task`, {
               task_id: taskId,
-              consecutive_failures: consecutiveEcsPollFailures,
+              consecutive_failures: consecutiveComputePollFailures,
               error: err instanceof Error ? err.message : String(err),
             });
-            await failTask(taskId, ddbState.lastStatus, `ECS poll failed ${consecutiveEcsPollFailures} consecutive times: ${err instanceof Error ? err.message : String(err)}`, task.user_id, true);
+            await failTask(taskId, ddbState.lastStatus, `${computeLabel} poll failed ${consecutiveComputePollFailures} consecutive times: ${err instanceof Error ? err.message : String(err)}`, task.user_id, true);
             return { attempts: ddbState.attempts, lastStatus: TaskStatus.FAILED };
           }
-          logger.warn('ECS pollSession check failed (non-fatal)', {
+          logger.warn(`${computeLabel} pollSession check failed (non-fatal)`, {
             task_id: taskId,
-            consecutive_failures: consecutiveEcsPollFailures,
+            consecutive_failures: consecutiveComputePollFailures,
             error: err instanceof Error ? err.message : String(err),
           });
         }
       }
 
-      return { ...ddbState, consecutiveEcsPollFailures, consecutiveEcsCompletedPolls };
+      return { ...ddbState, consecutiveComputePollFailures, consecutiveComputeCompletedPolls };
     },
     {
       initialState: { attempts: 0 },
diff --git a/cdk/src/handlers/shared/orchestrator.ts b/cdk/src/handlers/shared/orchestrator.ts
@@ -47,10 +47,10 @@ export interface PollState {
   readonly lastStatus?: TaskStatusType;
   /** True when the agent stopped sending heartbeats while still RUNNING (likely crash/OOM). */
   readonly sessionUnhealthy?: boolean;
-  /** Consecutive ECS poll failures — escalated to error after 3. */
-  readonly consecutiveEcsPollFailures?: number;
-  /** Consecutive polls where ECS reports completed but DDB is not terminal — escalated after 5. */
-  readonly consecutiveEcsCompletedPolls?: number;
+  /** Consecutive compute poll failures — escalated to error after 3. */
+  readonly consecutiveComputePollFailures?: number;
+  /** Consecutive polls where compute reports completed but DDB is not terminal — escalated after 5. */
+  readonly consecutiveComputeCompletedPolls?: number;
 }
 
 /** After RUNNING this long, we expect `agent_heartbeat_at` from the agent (if ever set). */
diff --git a/cdk/src/handlers/shared/strategies/ec2-strategy.ts b/cdk/src/handlers/shared/strategies/ec2-strategy.ts
@@ -105,69 +105,75 @@ export class Ec2ComputeStrategy implements ComputeStrategy {
       ],
     }));
 
-    // 4. Build the boot command (mirrors ECS strategy env vars and Python boot command)
-    const envExports = [
-      `export TASK_ID='${taskId}'`,
-      `export REPO_URL='${String(payload.repo_url ?? '')}'`,
-      ...(payload.prompt ? [`export TASK_DESCRIPTION='${String(payload.prompt).replace(/'/g, "'\\''")}'`] : []),
-      ...(payload.issue_number ? [`export ISSUE_NUMBER='${String(payload.issue_number)}'`] : []),
-      `export MAX_TURNS='${String(payload.max_turns ?? 100)}'`,
-      ...(payload.max_budget_usd !== undefined ? [`export MAX_BUDGET_USD='${String(payload.max_budget_usd)}'`] : []),
-      ...(blueprintConfig.model_id ? [`export ANTHROPIC_MODEL='${blueprintConfig.model_id}'`] : []),
-      ...(blueprintConfig.system_prompt_overrides ? [`export SYSTEM_PROMPT_OVERRIDES='${blueprintConfig.system_prompt_overrides.replace(/'/g, "'\\''")}'`] : []),
-      "export CLAUDE_CODE_USE_BEDROCK='1'",
-      ...(payload.github_token_secret_arn ? [`export GITHUB_TOKEN_SECRET_ARN='${String(payload.github_token_secret_arn)}'`] : []),
-      ...(payload.memory_id ? [`export MEMORY_ID='${String(payload.memory_id)}'`] : []),
-    ];
-
+    // 4. Build the boot script
+    // All task data is read from the S3 payload at runtime to avoid shell
+    // injection — no untrusted values are interpolated into the script.
+    // Only infrastructure constants (bucket name, ECR URI) are embedded.
     const bootScript = [
       '#!/bin/bash',
       'set -euo pipefail',
       '',
+      '# Derive region from IMDS (SSM does not always set AWS_REGION)',
+      'export AWS_REGION=$(ec2-metadata --availability-zone | cut -d" " -f2 | sed \'s/.$/\'\'/)\'',
+      'export AWS_DEFAULT_REGION="$AWS_REGION"',
+      '',
+      '# Resolve instance ID for tag cleanup',
+      'INSTANCE_ID=$(ec2-metadata -i | cut -d" " -f2)',
+      '',
+      '# Cleanup trap — always retag instance as idle on exit (success, error, or signal)',
+      'cleanup() {',
+      '  docker system prune -f || true',
+      '  rm -f /tmp/payload.json',
+      `  aws ec2 create-tags --resources "$INSTANCE_ID" --region "$AWS_REGION" --tags Key=bgagent:status,Value=idle || true`,
+      `  aws ec2 delete-tags --resources "$INSTANCE_ID" --region "$AWS_REGION" --tags Key=bgagent:task-id || true`,
+      '}',
+      'trap cleanup EXIT',
+      '',
       '# Fetch payload from S3',
       `aws s3 cp "s3://${EC2_PAYLOAD_BUCKET}/${payloadKey}" /tmp/payload.json`,
       'export AGENT_PAYLOAD=$(cat /tmp/payload.json)',
-      '',
-      '# Set environment variables',
-      ...envExports,
+      'export CLAUDE_CODE_USE_BEDROCK=1',
       '',
       '# ECR login and pull',
-      `aws ecr get-login-password --region $AWS_REGION | docker login --username AWS --password-stdin $(echo '${ECR_IMAGE_URI}' | cut -d/ -f1)`,
+      `aws ecr get-login-password --region "$AWS_REGION" | docker login --username AWS --password-stdin $(echo '${ECR_IMAGE_URI}' | cut -d/ -f1)`,
       `docker pull '${ECR_IMAGE_URI}'`,
       '',
-      '# Run the agent container',
-      'docker run --rm \\',
-      '  -e TASK_ID -e REPO_URL -e CLAUDE_CODE_USE_BEDROCK -e AGENT_PAYLOAD \\',
-      '  -e AWS_REGION -e AWS_DEFAULT_REGION \\',
-      `  ${payload.prompt ? '-e TASK_DESCRIPTION ' : ''}${payload.issue_number ? '-e ISSUE_NUMBER ' : ''}-e MAX_TURNS \\`,
-      `  ${payload.max_budget_usd !== undefined ? '-e MAX_BUDGET_USD ' : ''}${blueprintConfig.model_id ? '-e ANTHROPIC_MODEL ' : ''}${blueprintConfig.system_prompt_overrides ? '-e SYSTEM_PROMPT_OVERRIDES ' : ''}\\`,
-      `  ${payload.github_token_secret_arn ? '-e GITHUB_TOKEN_SECRET_ARN ' : ''}${payload.memory_id ? '-e MEMORY_ID ' : ''}\\`,
-      `  '${ECR_IMAGE_URI}' \\`,
+      '# Run the agent container — all config is read from AGENT_PAYLOAD inside the container',
+      `docker run --rm -e AGENT_PAYLOAD -e CLAUDE_CODE_USE_BEDROCK -e AWS_REGION -e AWS_DEFAULT_REGION '${ECR_IMAGE_URI}' \\`,
       '  python -c \'import json, os, sys; sys.path.insert(0, "/app"); from entrypoint import run_task; p = json.loads(os.environ["AGENT_PAYLOAD"]); r = run_task(repo_url=p.get("repo_url",""), task_description=p.get("prompt",""), issue_number=str(p.get("issue_number","")), github_token=p.get("github_token",""), anthropic_model=p.get("model_id",""), max_turns=int(p.get("max_turns",100)), max_budget_usd=p.get("max_budget_usd"), aws_region=os.environ.get("AWS_REGION",""), task_id=p.get("task_id",""), hydrated_context=p.get("hydrated_context"), system_prompt_overrides=p.get("system_prompt_overrides",""), prompt_version=p.get("prompt_version",""), memory_id=p.get("memory_id",""), task_type=p.get("task_type","new_task"), branch_name=p.get("branch_name",""), pr_number=str(p.get("pr_number",""))); sys.exit(0 if r.get("status")=="success" else 1)\'',
-      '',
-      '# Cleanup',
-      'docker system prune -f',
-      'rm -f /tmp/payload.json',
-      '',
-      '# Tag instance back to idle',
-      'INSTANCE_ID=$(ec2-metadata -i | cut -d" " -f2)',
-      'aws ec2 create-tags --resources "$INSTANCE_ID" --tags Key=bgagent:status,Value=idle',
-      'aws ec2 delete-tags --resources "$INSTANCE_ID" --tags Key=bgagent:task-id',
     ].join('\n');
 
-    // 5. Send SSM Run Command
-    const ssmResult = await getSsmClient().send(new SendCommandCommand({
-      DocumentName: 'AWS-RunShellScript',
-      InstanceIds: [instanceId],
-      Parameters: {
-        commands: [bootScript],
-      },
-      TimeoutSeconds: 32400, // 9 hours, matches orchestrator max
-    }));
+    // 5. Send SSM Run Command — rollback instance tags on failure
+    let commandId: string;
+    try {
+      const ssmResult = await getSsmClient().send(new SendCommandCommand({
+        DocumentName: 'AWS-RunShellScript',
+        InstanceIds: [instanceId],
+        Parameters: {
+          commands: [bootScript],
+        },
+        TimeoutSeconds: 32400, // 9 hours, matches orchestrator max
+      }));
 
-    const commandId = ssmResult.Command?.CommandId;
-    if (!commandId) {
-      throw new Error('SSM SendCommand returned no CommandId');
+      if (!ssmResult.Command?.CommandId) {
+        throw new Error('SSM SendCommand returned no CommandId');
+      }
+      commandId = ssmResult.Command.CommandId;
+    } catch (err) {
+      // Rollback: retag instance as idle so it's not stuck as busy
+      try {
+        await getEc2Client().send(new CreateTagsCommand({
+          Resources: [instanceId],
+          Tags: [{ Key: 'bgagent:status', Value: 'idle' }],
+        }));
+        await getEc2Client().send(new DeleteTagsCommand({
+          Resources: [instanceId],
+          Tags: [{ Key: 'bgagent:task-id' }],
+        }));
+      } catch {
+        logger.warn('Failed to rollback instance tags after dispatch failure', { instance_id: instanceId, task_id: taskId });
+      }
+      throw err;
     }
 
     logger.info('EC2 SSM command dispatched', {
diff --git a/cdk/src/stacks/agent.ts b/cdk/src/stacks/agent.ts
@@ -339,7 +339,6 @@ export class AgentStack extends Stack {
       //   fleetTagValue: ec2Fleet.fleetTagValue,
       //   payloadBucketName: ec2Fleet.payloadBucket.bucketName,
       //   ecrImageUri: agentImageAsset.imageUri,
-      //   instanceRoleArn: ec2Fleet.instanceRole.roleArn,
       // },
     });
 
@@ -366,7 +365,7 @@ export class AgentStack extends Stack {
       // To allow cancel-task to stop ECS-backed tasks, uncomment:
       // ecsClusterArn: ecsCluster.cluster.clusterArn,
       // To allow cancel-task to stop EC2-backed tasks, uncomment:
-      // ec2FleetConfig: { instanceRoleArn: ec2Fleet.instanceRole.roleArn },
+      // ec2FleetConfig: {},
     });
 
     // --- Operator dashboard ---
