Skip to content

Commit cee7331

Browse files
MichaelWalker-gitbgagent
MichaelWalker-git
and
bgagent
authored
feat(compute): add ComputeStrategy interface with ECS Fargate backend (#8)
* feat(orchestrator): extract ComputeStrategy interface from hardcoded AgentCore logic Introduce ComputeStrategy interface with SessionHandle/SessionStatus types and resolveComputeStrategy factory. Extract AgentCoreComputeStrategy from orchestrator.ts. Refactor orchestrate-task handler to use strategy pattern for session lifecycle (start/poll/stop). Pure refactor — no behavior change, identical CloudFormation output. * fix(ci): pass GITHUB_TOKEN to mise to avoid API rate limits The mise install step downloads tools (trivy) from GitHub releases. Without GITHUB_TOKEN, unauthenticated requests hit the 60 req/hr rate limit, causing flaky CI failures. * fix(ci): set GITHUB_API_TOKEN for mise tool downloads Mise uses GITHUB_API_TOKEN (not GITHUB_TOKEN) for authenticated GitHub API requests when downloading aqua tools like trivy. * fix(ci): disable security-only tools in build workflow Trivy, grype, semgrep, osv-scanner, and gitleaks are only needed for security scanning tasks, not for the build/test/synth pipeline. Disable them via MISE_DISABLE_TOOLS to avoid GitHub API rate limits when mise tries to download them on every PR build. * fix: address PR review comments - Keep gitleaks and osv-scanner enabled in CI build (only disable trivy/grype/semgrep which need GitHub API downloads) - Remove unused @aws-sdk/client-bedrock-agentcore mock from orchestrate-task.test.ts (SDK is no longer imported by orchestrator) - Update PR description to note additive strategy_type event field * fix: address Alain's PR review findings 1. Single source of truth for runtimeArn — removed constructor param, strategy now reads exclusively from blueprintConfig.runtime_arn 2. Lazy singleton for BedrockAgentCoreClient — module-level shared client avoids creating new TLS sessions per invocation 3. ComputeType union type ('agentcore' | 'ecs') with exhaustive switch and never-pattern in resolveComputeStrategy 4. Differentiated error handling in stopSession — ResourceNotFoundException (info), ThrottlingException/AccessDeniedException (error), others (warn) 5. Added logger.info('Session started') after full invoke+transition+event sequence in orchestrate-task.ts 6. Added start-session-composition.test.ts with integration tests for happy path, error path (failTask), and partial failure (transitionTask throws) 7. pollSession now throws NotImplementedError instead of returning stale 'running' status — clear signal for future developers * fix: resolve ESLint errors in test files - Replace require() with ES import for BedrockAgentCoreClient mock - Fix import ordering in start-session-composition test * feat(compute): implement ECS Fargate backend via ComputeStrategy pattern Wire ECS Fargate as a compute backend behind the existing ComputeStrategy interface, using the existing durable Lambda orchestrator. No separate stacks or Step Functions — ECS is a strategy option alongside AgentCore. Changes: - EcsComputeStrategy: startSession (RunTask), pollSession (DescribeTasks state mapping), stopSession (StopTask with graceful error handling) - EcsAgentCluster construct: ECS Cluster (container insights), Fargate task def (2 vCPU/4GB/ARM64), security group (TCP 443 egress only), CloudWatch log group, task role (DynamoDB, SecretsManager, Bedrock) - TaskOrchestrator: optional ECS props for env vars and IAM policies (ecs:RunTask/DescribeTasks/StopTask conditioned on cluster ARN, iam:PassRole conditioned on ecs-tasks.amazonaws.com) - Orchestrator polling: ECS compute-level crash detection alongside existing DDB polling (non-fatal, wrapped in try/catch) - AgentStack: conditional ECS infrastructure (ABCA_ENABLE_ECS env var) - Full test coverage: 15 ECS strategy tests, 9 construct tests, 5 orchestrator ECS tests. All 563 tests pass. Deployed and verified: stack deploys cleanly, CDK synth passes cdk-nag, agent task running on AgentCore path unaffected. * fix: address Copilot PR review comments on PR #8 - Keep gitleaks/osv-scanner enabled in CI (only disable trivy/grype/semgrep) - Type ComputeStrategy.type and SessionHandle.strategyType as ComputeType - Trim/filter ECS_SUBNETS to handle whitespace and trailing commas - Handle undefined exit code in ECS pollSession (container never started) - Scope iam:PassRole to specific ECS task/execution role ARNs - Validate all-or-nothing ECS props in TaskOrchestrator constructor - Remove dead hasEcsBlueprint detection; document env-flag driven approach - Add comment noting strategy_type as additive event field * fix: remove duplicate ECS validation block after rebase * fix: remove local-docs from repo and add to .gitignore * fix(ecs): override container command to run agent in batch mode The ECS container's default CMD starts uvicorn server:app which waits for HTTP POST to /invocations — but in standalone ECS nobody sends that request, leaving the agent idle. Override the container command to invoke entrypoint.run_task() directly with the full orchestrator payload via AGENT_PAYLOAD env var. Also add GITHUB_TOKEN_SECRET_ARN to the ECS task definition base environment. * chore(pr): address review * chore(project): update docs --------- Co-authored-by: bgagent <bgagent@noreply.github.com>
1 parent b9c9ecc commit cee7331

28 files changed

Lines changed: 2690 additions & 100 deletions

‎.github/workflows/build.yml‎

Lines changed: 6 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎.gitignore‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎cdk/package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
"@aws-cdk/mixins-preview": "2.238.0-alpha.0",
1919
"@aws-sdk/client-bedrock-agentcore": "^3.1021.0",
2020
"@aws-sdk/client-bedrock-runtime": "^3.1021.0",
21+
"@aws-sdk/client-ecs": "^3.1021.0",
2122
"@aws-sdk/client-dynamodb": "^3.1021.0",
2223
"@aws-sdk/client-lambda": "^3.1021.0",
2324
"@aws-sdk/client-secrets-manager": "^3.1021.0",

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

Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,156 @@
1+
/**
2+
* MIT No Attribution
3+
*
4+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
5+
*
6+
* Permission is hereby granted, free of charge, to any person obtaining a copy of
7+
* the Software without restriction, including without limitation the rights to
8+
* use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
9+
* the Software, and to permit persons to whom the Software is furnished to do so.
10+
*
11+
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
12+
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
13+
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
14+
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
15+
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
16+
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
17+
* SOFTWARE.
18+
*/
19+
20+
import { RemovalPolicy } from 'aws-cdk-lib';
21+
import * as dynamodb from 'aws-cdk-lib/aws-dynamodb';
22+
import * as ec2 from 'aws-cdk-lib/aws-ec2';
23+
import * as ecr_assets from 'aws-cdk-lib/aws-ecr-assets';
24+
import * as ecs from 'aws-cdk-lib/aws-ecs';
25+
import * as iam from 'aws-cdk-lib/aws-iam';
26+
import * as logs from 'aws-cdk-lib/aws-logs';
27+
import * as secretsmanager from 'aws-cdk-lib/aws-secretsmanager';
28+
import { NagSuppressions } from 'cdk-nag';
29+
import { Construct } from 'constructs';
30+
31+
export interface EcsAgentClusterProps {
32+
readonly vpc: ec2.IVpc;
33+
readonly agentImageAsset: ecr_assets.DockerImageAsset;
34+
readonly taskTable: dynamodb.ITable;
35+
readonly taskEventsTable: dynamodb.ITable;
36+
readonly userConcurrencyTable: dynamodb.ITable;
37+
readonly githubTokenSecret: secretsmanager.ISecret;
38+
readonly memoryId?: string;
39+
}
40+
41+
export class EcsAgentCluster extends Construct {
42+
public readonly cluster: ecs.Cluster;
43+
public readonly taskDefinition: ecs.FargateTaskDefinition;
44+
public readonly securityGroup: ec2.SecurityGroup;
45+
public readonly containerName: string;
46+
public readonly taskRoleArn: string;
47+
public readonly executionRoleArn: string;
48+
49+
constructor(scope: Construct, id: string, props: EcsAgentClusterProps) {
50+
super(scope, id);
51+
52+
this.containerName = 'AgentContainer';
53+
54+
// ECS Cluster with Fargate capacity provider and container insights
55+
this.cluster = new ecs.Cluster(this, 'Cluster', {
56+
vpc: props.vpc,
57+
containerInsights: true,
58+
});
59+
60+
// Security group — egress TCP 443 only
61+
this.securityGroup = new ec2.SecurityGroup(this, 'TaskSG', {
62+
vpc: props.vpc,
63+
description: 'ECS Agent Tasks - egress TCP 443 only',
64+
allowAllOutbound: false,
65+
});
66+
67+
this.securityGroup.addEgressRule(
68+
ec2.Peer.anyIpv4(),
69+
ec2.Port.tcp(443),
70+
'Allow HTTPS egress (GitHub API, AWS services)',
71+
);
72+
73+
// CloudWatch log group for agent task output
74+
const logGroup = new logs.LogGroup(this, 'TaskLogGroup', {
75+
retention: logs.RetentionDays.THREE_MONTHS,
76+
removalPolicy: RemovalPolicy.DESTROY,
77+
});
78+
79+
// Task execution role (used by ECS agent to pull images, write logs)
80+
// CDK creates this automatically via taskDefinition, but we need to
81+
// grant additional permissions to the task role.
82+
83+
// Fargate task definition
84+
this.taskDefinition = new ecs.FargateTaskDefinition(this, 'TaskDef', {
85+
cpu: 2048,
86+
memoryLimitMiB: 4096,
87+
runtimePlatform: {
88+
cpuArchitecture: ecs.CpuArchitecture.ARM64,
89+
operatingSystemFamily: ecs.OperatingSystemFamily.LINUX,
90+
},
91+
});
92+
93+
// Container
94+
this.taskDefinition.addContainer(this.containerName, {
95+
image: ecs.ContainerImage.fromDockerImageAsset(props.agentImageAsset),
96+
logging: ecs.LogDrivers.awsLogs({
97+
logGroup,
98+
streamPrefix: 'agent',
99+
}),
100+
environment: {
101+
CLAUDE_CODE_USE_BEDROCK: '1',
102+
TASK_TABLE_NAME: props.taskTable.tableName,
103+
TASK_EVENTS_TABLE_NAME: props.taskEventsTable.tableName,
104+
USER_CONCURRENCY_TABLE_NAME: props.userConcurrencyTable.tableName,
105+
LOG_GROUP_NAME: logGroup.logGroupName,
106+
GITHUB_TOKEN_SECRET_ARN: props.githubTokenSecret.secretArn,
107+
...(props.memoryId && { MEMORY_ID: props.memoryId }),
108+
},
109+
});
110+
111+
// Task role permissions
112+
const taskRole = this.taskDefinition.taskRole;
113+
114+
// DynamoDB read/write on task tables
115+
props.taskTable.grantReadWriteData(taskRole);
116+
props.taskEventsTable.grantReadWriteData(taskRole);
117+
props.userConcurrencyTable.grantReadWriteData(taskRole);
118+
119+
// Secrets Manager read for GitHub token
120+
props.githubTokenSecret.grantRead(taskRole);
121+
122+
// Bedrock model invocation
123+
taskRole.addToPrincipalPolicy(new iam.PolicyStatement({
124+
actions: [
125+
'bedrock:InvokeModel',
126+
'bedrock:InvokeModelWithResponseStream',
127+
],
128+
resources: ['*'],
129+
}));
130+
131+
// CloudWatch Logs write
132+
logGroup.grantWrite(taskRole);
133+
134+
// Expose role ARNs for scoped iam:PassRole in the orchestrator
135+
this.taskRoleArn = taskRole.roleArn;
136+
this.executionRoleArn = this.taskDefinition.executionRole!.roleArn;
137+
138+
NagSuppressions.addResourceSuppressions(this.taskDefinition, [
139+
{
140+
id: 'AwsSolutions-IAM5',
141+
reason: 'DynamoDB index/* wildcards generated by CDK grantReadWriteData; Bedrock InvokeModel requires * resource; Secrets Manager wildcards from CDK grantRead; CloudWatch Logs wildcards from CDK grantWrite',
142+
},
143+
{
144+
id: 'AwsSolutions-ECS2',
145+
reason: 'Environment variables contain table names and configuration, not secrets — GitHub token is fetched from Secrets Manager at runtime',
146+
},
147+
], true);
148+
149+
NagSuppressions.addResourceSuppressions(this.cluster, [
150+
{
151+
id: 'AwsSolutions-ECS4',
152+
reason: 'Container insights is enabled via the containerInsights prop',
153+
},
154+
], true);
155+
}
156+
}

‎cdk/src/constructs/task-api.ts‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,12 @@ export interface TaskApiProps {
100100
* First ARN is also passed as `RUNTIME_ARN` when the task record has no `agent_runtime_arn`.
101101
*/
102102
readonly agentCoreStopSessionRuntimeArns?: string[];
103+
104+
/**
105+
* ECS cluster ARN for cancel-task to stop ECS-backed tasks.
106+
* When provided, the cancel Lambda gets `ECS_CLUSTER_ARN` env var and `ecs:StopTask` permission.
107+
*/
108+
readonly ecsClusterArn?: string;
103109
}
104110

105111
/**
@@ -329,6 +335,9 @@ export class TaskApi extends Construct {
329335
if (stopSessionArns.length > 0) {
330336
cancelTaskEnv.RUNTIME_ARN = stopSessionArns[0]!;
331337
}
338+
if (props.ecsClusterArn) {
339+
cancelTaskEnv.ECS_CLUSTER_ARN = props.ecsClusterArn;
340+
}
332341

333342
const cancelTaskFn = new lambda.NodejsFunction(this, 'CancelTaskFn', {
334343
entry: path.join(handlersDir, 'cancel-task.ts'),
@@ -363,6 +372,18 @@ export class TaskApi extends Construct {
363372
}));
364373
}
365374

375+
if (props.ecsClusterArn) {
376+
cancelTaskFn.addToRolePolicy(new iam.PolicyStatement({
377+
actions: ['ecs:StopTask'],
378+
resources: ['*'],
379+
conditions: {
380+
ArnEquals: {
381+
'ecs:cluster': props.ecsClusterArn,
382+
},
383+
},
384+
}));
385+
}
386+
366387
// Repo table read for onboarding gate
367388
if (props.repoTable) {
368389
props.repoTable.grantReadData(createTaskFn);

‎cdk/src/constructs/task-orchestrator.ts‎

Lines changed: 50 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,21 @@ export interface TaskOrchestratorProps {
112112
* Bedrock Guardrail version. Required when guardrailId is provided.
113113
*/
114114
readonly guardrailVersion?: string;
115+
116+
/**
117+
* ECS Fargate compute strategy configuration.
118+
* When provided, ECS-related env vars and IAM policies are added to the orchestrator.
119+
* All fields are required — this makes the all-or-nothing constraint self-evident at the type level.
120+
*/
121+
readonly ecsConfig?: {
122+
readonly clusterArn: string;
123+
readonly taskDefinitionArn: string;
124+
readonly subnets: string;
125+
readonly securityGroup: string;
126+
readonly containerName: string;
127+
readonly taskRoleArn: string;
128+
readonly executionRoleArn: string;
129+
};
115130
}
116131

117132
/**
@@ -173,6 +188,13 @@ export class TaskOrchestrator extends Construct {
173188
...(props.memoryId && { MEMORY_ID: props.memoryId }),
174189
...(props.guardrailId && { GUARDRAIL_ID: props.guardrailId }),
175190
...(props.guardrailVersion && { GUARDRAIL_VERSION: props.guardrailVersion }),
191+
...(props.ecsConfig && {
192+
ECS_CLUSTER_ARN: props.ecsConfig.clusterArn,
193+
ECS_TASK_DEFINITION_ARN: props.ecsConfig.taskDefinitionArn,
194+
ECS_SUBNETS: props.ecsConfig.subnets,
195+
ECS_SECURITY_GROUP: props.ecsConfig.securityGroup,
196+
ECS_CONTAINER_NAME: props.ecsConfig.containerName,
197+
}),
176198
},
177199
bundling: {
178200
externalModules: ['@aws-sdk/*'],
@@ -213,6 +235,33 @@ export class TaskOrchestrator extends Construct {
213235
resources: runtimeResources,
214236
}));
215237

238+
// ECS compute strategy permissions (only when ECS is configured)
239+
if (props.ecsConfig) {
240+
this.fn.addToRolePolicy(new iam.PolicyStatement({
241+
actions: [
242+
'ecs:RunTask',
243+
'ecs:DescribeTasks',
244+
'ecs:StopTask',
245+
],
246+
resources: ['*'],
247+
conditions: {
248+
ArnEquals: {
249+
'ecs:cluster': props.ecsConfig.clusterArn,
250+
},
251+
},
252+
}));
253+
254+
this.fn.addToRolePolicy(new iam.PolicyStatement({
255+
actions: ['iam:PassRole'],
256+
resources: [props.ecsConfig.taskRoleArn, props.ecsConfig.executionRoleArn],
257+
conditions: {
258+
StringEquals: {
259+
'iam:PassedToService': 'ecs-tasks.amazonaws.com',
260+
},
261+
},
262+
}));
263+
}
264+
216265
// Per-repo Secrets Manager grants (e.g. per-repo GitHub tokens from Blueprints)
217266
for (const [index, secretArn] of (props.additionalSecretArns ?? []).entries()) {
218267
const secret = secretsmanager.Secret.fromSecretCompleteArn(
@@ -264,7 +313,7 @@ export class TaskOrchestrator extends Construct {
264313
},
265314
{
266315
id: 'AwsSolutions-IAM5',
267-
reason: 'DynamoDB index/* wildcards generated by CDK grantReadWriteData; AgentCore runtime/* required for sub-resource invocation; Secrets Manager wildcards generated by CDK grantRead; AgentCore Memory wildcards generated by CDK grantRead/grantWrite',
316+
reason: 'DynamoDB index/* wildcards generated by CDK grantReadWriteData; AgentCore runtime/* required for sub-resource invocation; Secrets Manager wildcards generated by CDK grantRead; AgentCore Memory wildcards generated by CDK grantRead/grantWrite; ECS RunTask/DescribeTasks/StopTask conditioned on cluster ARN; iam:PassRole scoped to ECS task/execution roles and conditioned on ecs-tasks.amazonaws.com',
268317
},
269318
], true);
270319
}

‎cdk/src/handlers/cancel-task.ts‎

Lines changed: 52 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
import { BedrockAgentCoreClient, StopRuntimeSessionCommand } from '@aws-sdk/client-bedrock-agentcore';
2121
import { DynamoDBClient } from '@aws-sdk/client-dynamodb';
22+
import { ECSClient, StopTaskCommand } from '@aws-sdk/client-ecs';
2223
import { DynamoDBDocumentClient, GetCommand, PutCommand, UpdateCommand } from '@aws-sdk/lib-dynamodb';
2324
import type { APIGatewayProxyEvent, APIGatewayProxyResult } from 'aws-lambda';
2425
import { ulid } from 'ulid';
@@ -31,10 +32,12 @@ import { computeTtlEpoch } from './shared/validation';
3132

3233
const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({}));
3334
const agentCoreClient = new BedrockAgentCoreClient({});
35+
const ecsClient = new ECSClient({});
3436
const TABLE_NAME = process.env.TASK_TABLE_NAME!;
3537
const EVENTS_TABLE_NAME = process.env.TASK_EVENTS_TABLE_NAME!;
3638
const TASK_RETENTION_DAYS = Number(process.env.TASK_RETENTION_DAYS ?? '90');
3739
const RUNTIME_ARN = process.env.RUNTIME_ARN;
40+
const ECS_CLUSTER_ARN = process.env.ECS_CLUSTER_ARN;
3841

3942
/**
4043
* DELETE /v1/tasks/{task_id} — Cancel a task.
@@ -107,19 +110,57 @@ export async function handler(event: APIGatewayProxyEvent): Promise<APIGatewayPr
107110
throw condErr;
108111
}
109112

110-
// 6b. Stop AgentCore runtime session so the container winds down (best-effort)
111-
if (wasRunning && runtimeSessionId && agentRuntimeArn) {
112-
try {
113-
await agentCoreClient.send(new StopRuntimeSessionCommand({
114-
runtimeSessionId: runtimeSessionId,
115-
agentRuntimeArn: agentRuntimeArn,
116-
}));
117-
logger.info('StopRuntimeSession invoked after cancel', { task_id: taskId, request_id: requestId });
118-
} catch (stopErr) {
119-
logger.warn('StopRuntimeSession failed after cancel (session may already be gone)', {
113+
// 6b. Stop the compute session so the container winds down (best-effort)
114+
if (wasRunning && runtimeSessionId) {
115+
const computeType = record.compute_type;
116+
if (computeType === 'ecs') {
117+
// ECS-backed task — stop the Fargate task
118+
const clusterArn = record.compute_metadata?.clusterArn ?? ECS_CLUSTER_ARN;
119+
const taskArn = record.compute_metadata?.taskArn;
120+
if (clusterArn && taskArn) {
121+
try {
122+
await ecsClient.send(new StopTaskCommand({
123+
cluster: clusterArn,
124+
task: taskArn,
125+
reason: 'Cancelled by user',
126+
}));
127+
logger.info('ECS StopTask invoked after cancel', { task_id: taskId, ecs_task_arn: taskArn, request_id: requestId });
128+
} catch (stopErr) {
129+
logger.warn('ECS StopTask failed after cancel (task may already be stopped)', {
130+
task_id: taskId,
131+
request_id: requestId,
132+
error: stopErr instanceof Error ? stopErr.message : String(stopErr),
133+
});
134+
}
135+
} else {
136+
logger.warn('ECS task cancel skipped: missing clusterArn or taskArn in compute_metadata', {
137+
task_id: taskId,
138+
request_id: requestId,
139+
has_cluster: !!clusterArn,
140+
has_task: !!taskArn,
141+
});
142+
}
143+
} else if (agentRuntimeArn) {
144+
// AgentCore-backed task (default)
145+
try {
146+
await agentCoreClient.send(new StopRuntimeSessionCommand({
147+
runtimeSessionId: runtimeSessionId,
148+
agentRuntimeArn: agentRuntimeArn,
149+
}));
150+
logger.info('StopRuntimeSession invoked after cancel', { task_id: taskId, request_id: requestId });
151+
} catch (stopErr) {
152+
logger.warn('StopRuntimeSession failed after cancel (session may already be gone)', {
153+
task_id: taskId,
154+
request_id: requestId,
155+
error: stopErr instanceof Error ? stopErr.message : String(stopErr),
156+
});
157+
}
158+
} else {
159+
logger.warn('Running task has no recognized compute backend to stop', {
120160
task_id: taskId,
121161
request_id: requestId,
122-
error: stopErr instanceof Error ? stopErr.message : String(stopErr),
162+
compute_type: computeType,
163+
has_runtime_arn: !!agentRuntimeArn,
123164
});
124165
}
125166
}

0 commit comments

Comments
 (0)