Building a Serverless AI Training Platform on AWS: Pay-as-you-Grow Architecture with Bedrock, SQS, and WebSockets

[Browser — Vite + React SPA]
        ↓ CloudFront + S3 (static hosting)
[Edge: Cognito · API Gateway REST · API Gateway WebSocket]
        ↓ JWT-validated
[Compute: Lambda functions (Node.js 22 / TypeScript)]
   ├── Mission Generator  →  SQS  →  AI Worker  →  Bedrock (non-streaming)
   ├── Chat Hints                 →  Bedrock (streaming)
   ├── History & Profile          →  DynamoDB CRUD
   ├── Auto-Grader                →  template parse + criteria check
   └── WS Connection Manager      →  ws_connections table
        ↓
[Data: DynamoDB — users · missions · conversations · ws_connections]

[Client] → API Gateway → Lambda → Bedrock.invokeModel → Lambda response → Client

- Client → API Gateway → Lambda → Bedrock → Client
+ Client → API Gateway → MissionGenerator(Lambda) → SQS
+                                                   ↓
+                                                AIWorker(Lambda) → Bedrock
+                                                   ↓
+                                        DynamoDB.update + WS.push("mission ready")
+                                                   ↓
+                                                Client (via WebSocket)

export const handler = async (event: SQSEvent) => {
  for (const record of event.Records) {
    const { missionId, userId, prompt } = JSON.parse(record.body);
    try {
      const result = await invokeBedrock(prompt);
      await updateMission(missionId, { status: "ready", result });
      await pushToUser(userId, { type: "mission_ready", missionId });
    } catch (err) {
      if (isThrottling(err)) throw err;            // back to SQS
      await updateMission(missionId, { status: "failed", error: String(err) });
    }
  }
};

const ticket = event.queryStringParameters?.ticket;
const row = await ddb.send(new DeleteItemCommand({
  TableName: TICKETS_TABLE,
  Key: { ticket: { S: ticket } },
  ReturnValues: "ALL_OLD",                          // atomic consume
}));
if (!row.Attributes) return deny();
return allow(row.Attributes.userId.S);

const MODEL_ID = "global.anthropic.claude-sonnet-4-6";
// NOT "anthropic.claude-sonnet-4-6-20250101-v1:0"

{
  "resourceType": "AWS::Lambda::Function",
  "property": "Runtime",
  "expected": "nodejs22.x",
  "label": "Lambda uses Node.js 22 runtime"
}

[Client uploads template.yaml/.json]
        ↓
[POST /missions/{id}/grade] → AutoGrader Lambda
        ↓
1. Auth check (Cognito JWT) — verify user owns this mission
2. Load mission from DynamoDB → get successCriteria[]
3. Parse template body:
     • try JSON.parse first
     • fall back to js-yaml with custom !Ref/!GetAtt/!Sub tags
4. For each criterion: walk template.Resources, check property path
5. score = passed / total
6. Persist {gradeStatus, gradeResult} back to mission row
7. Return JSON

// 1. JSON-or-YAML parser, with custom tags so !Ref etc. don't crash
const CFN_TAG_NAMES = ['Ref', 'GetAtt', 'Sub', 'Select', 'If', /* ... */];
const cfnTypes = CFN_TAG_NAMES.flatMap(tag =>
  (['scalar', 'sequence', 'mapping'] as const).map(kind =>
    new Type(`!${tag}`, { kind, construct: (d: unknown) => d })));
const CFN_SCHEMA = DEFAULT_SCHEMA.extend(cfnTypes);

// 2. Walks dotted paths like "LambdaConfigurations[0].Event"
function getNestedValue(obj: Record<string, unknown>, path: string): unknown {
  return path.split('.').reduce<unknown>((acc, key) => {
    if (acc == null || typeof acc !== 'object') return undefined;
    const arrMatch = key.match(/^(.+?)\[(\d+)\]$/);
    if (arrMatch) {
      const arr = (acc as Record<string, unknown>)[arrMatch[1]];
      return Array.isArray(arr) ? arr[Number(arrMatch[2])] : undefined;
    }
    return (acc as Record<string, unknown>)[key];
  }, obj);
}

// 3. Exact match, escalates to regex if expected has wildcards
function matchesValue(val: unknown, expected: string): boolean {
  const str = String(val);
  if (!expected.includes('*') && !expected.includes('?')) return str === expected;
  const pattern = expected
    .replace(/[.+^${}()|[\]\\]/g, '\\$&')
    .replace(/\*/g, '.*').replace(/\?/g, '.');
  return new RegExp(`^${pattern}$`).test(str);
}

this.missionsTable = new dynamodb.Table(this, 'Missions', {
  partitionKey: { name: 'missionId', type: dynamodb.AttributeType.STRING },
  billingMode: dynamodb.BillingMode.PAY_PER_REQUEST,    // on-demand → zero idle cost
  removalPolicy: cdk.RemovalPolicy.RETAIN,              // never auto-delete user data
});
this.missionsTable.addGlobalSecondaryIndex({
  indexName: 'userId-createdAt-index',
  partitionKey: { name: 'userId', type: dynamodb.AttributeType.STRING },
  sortKey: { name: 'createdAt', type: dynamodb.AttributeType.STRING },
});

this.wsConnectionsTable = new dynamodb.Table(this, 'WsConnections', {
  partitionKey: { name: 'connectionId', type: dynamodb.AttributeType.STRING },
  billingMode: dynamodb.BillingMode.PAY_PER_REQUEST,
  timeToLiveAttribute: 'ttl',                           // DDB sweeps stale connections
  removalPolicy: cdk.RemovalPolicy.DESTROY,             // ephemeral state, safe to drop
});

this.wsTicketsTable = new dynamodb.Table(this, 'WsTickets', {
  partitionKey: { name: 'ticket', type: dynamodb.AttributeType.STRING },
  billingMode: dynamodb.BillingMode.PAY_PER_REQUEST,
  timeToLiveAttribute: 'ttl',                           // 60s TTL — defence in depth
  removalPolicy: cdk.RemovalPolicy.DESTROY,
});

const dlq = new sqs.Queue(this, 'Dlq', {
  retentionPeriod: cdk.Duration.days(14),
});

new cloudwatch.Alarm(this, 'DlqDepthAlarm', {
  metric: dlq.metricApproximateNumberOfMessagesVisible(),
  threshold: 1,
  evaluationPeriods: 1,
  comparisonOperator: cloudwatch.ComparisonOperator.GREATER_THAN_OR_EQUAL_TO_THRESHOLD,
  alarmDescription: 'Mission generation DLQ has messages — check for stuck missions',
  treatMissingData: cloudwatch.TreatMissingData.NOT_BREACHING,
});

const missionQueue = new sqs.Queue(this, 'MissionQueue', {
  visibilityTimeout: cdk.Duration.seconds(1800),        // 6× Lambda timeout
  deadLetterQueue: { queue: dlq, maxReceiveCount: 3 },
});

const sharedEnv = {
  USERS_TABLE: props.usersTable.tableName,
  MISSIONS_TABLE: props.missionsTable.tableName,
  WS_ENDPOINT: wsStage.callbackUrl,
  MISSION_QUEUE_URL: missionQueue.queueUrl,
  BEDROCK_MODEL_ID: 'global.anthropic.claude-sonnet-4-6',
  // ...
};

const mkFn = (id: string, handler: string, extra?: Partial<lambda.FunctionProps>) =>
  new lambda.Function(this, id, {
    runtime: lambda.Runtime.NODEJS_22_X,
    handler,
    code: lambda.Code.fromAsset('../backend/dist'),
    timeout: cdk.Duration.minutes(5),
    environment: sharedEnv,
    logGroup: new logs.LogGroup(this, `${id}Logs`, {
      retention: logs.RetentionDays.ONE_MONTH,
      removalPolicy: cdk.RemovalPolicy.DESTROY,
    }),
    ...extra,
  });

const missionGeneratorFn = mkFn('MissionGenerator', 'handlers/mission-generator.handler');
const aiWorkerFn = mkFn('AiWorker', 'handlers/ai-worker.handler', {
  reservedConcurrentExecutions: 10,                     // cap Bedrock fan-out
});

// Wire the AI worker to consume from SQS
missionQueue.grantSendMessages(missionGeneratorFn);
missionQueue.grantConsumeMessages(aiWorkerFn);
aiWorkerFn.addEventSource(new lambdaEventSources.SqsEventSource(missionQueue, {
  batchSize: 1,                                         // one mission per invocation
}));

// Per-function, per-table DynamoDB grants
props.missionsTable.grantWriteData(missionGeneratorFn);
props.missionsTable.grantReadWriteData(aiWorkerFn);
props.missionsTable.grantReadData(chatHintsFn);

// Bedrock + WebSocket management policies for the AI-touching functions
const bedrockPolicy = new iam.PolicyStatement({
  actions: ['bedrock:InvokeModel', 'bedrock:InvokeModelWithResponseStream'],
  resources: ['*'],
});
const wsManagePolicy = new iam.PolicyStatement({
  actions: ['execute-api:ManageConnections'],
  resources: [`arn:aws:execute-api:${this.region}:${this.account}:${wsApi.apiId}/*`],
});
[aiWorkerFn, chatHintsFn].forEach(f => {
  f.addToRolePolicy(bedrockPolicy);
  f.addToRolePolicy(wsManagePolicy);
});

// WebSocket API (must exist before Lambdas, so callbackUrl can be injected as env)
const wsApi = new apigwv2.WebSocketApi(this, 'WsApi');
const wsStage = new apigwv2.WebSocketStage(this, 'WsStage', {
  webSocketApi: wsApi, stageName: 'prod', autoDeploy: true,
});

const wsi = (fn: lambda.Function, id: string) =>
  new integrations.WebSocketLambdaIntegration(id, fn);
wsApi.addRoute('$connect',    { integration: wsi(wsConnectionFn, 'Connect') });
wsApi.addRoute('$disconnect', { integration: wsi(wsConnectionFn, 'Disconnect') });
wsApi.addRoute('sendHint',    { integration: wsi(chatHintsFn, 'SendHint') });

// REST API with Cognito authorizer
const rest = new apigw.RestApi(this, 'RestApi', {
  deployOptions: { stageName: 'prod' },
  defaultCorsPreflightOptions: {
    allowOrigins: apigw.Cors.ALL_ORIGINS,
    allowHeaders: apigw.Cors.DEFAULT_HEADERS,
    allowMethods: apigw.Cors.ALL_METHODS,
  },
});
const cogAuth = new apigw.CognitoUserPoolsAuthorizer(this, 'CogAuth', {
  cognitoUserPools: [props.userPool],
});
const opts = { authorizer: cogAuth, authorizationType: apigw.AuthorizationType.COGNITO };

const missionsRes = rest.root.addResource('missions');
missionsRes.addMethod('POST', new apigw.LambdaIntegration(missionGeneratorFn), opts);
missionsRes.addResource('{id}').addResource('grade')
  .addMethod('POST', new apigw.LambdaIntegration(autoGraderFn), opts);

const app = new cdk.App();
const env = { region: 'ap-southeast-1' };

const db   = new DatabaseStack(app, 'AwsMissionDatabase', { env });
const auth = new AuthStack(app, 'AwsMissionAuth', { env });
const store = new StorageStack(app, 'AwsMissionStorage', { env });
new BackendStack(app, 'AwsMissionBackend', {
  env,
  ...db, ...auth,                                       // pass tables + user pool as props
});