feat(orchestration): admission queue with deferred pickup (#441)

cdk/src/constructs/admission-queue-pickup.ts

@@ -0,0 +1,160 @@
+/**
+ *  MIT No Attribution
+ *
+ *  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *  Permission is hereby granted, free of charge, to any person obtaining a copy of
+ *  the Software without restriction, including without limitation the rights to
+ *  use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
+ *  the Software, and to permit persons to whom the Software is furnished to do so.
+ *
+ *  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ *  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ *  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ *  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ *  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ *  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ *  SOFTWARE.
+ */
+
+import * as path from 'path';
+import { Duration } from 'aws-cdk-lib';
+import * as dynamodb from 'aws-cdk-lib/aws-dynamodb';
+import * as events from 'aws-cdk-lib/aws-events';
+import * as targets from 'aws-cdk-lib/aws-events-targets';
+import * as iam from 'aws-cdk-lib/aws-iam';
+import { Architecture, Runtime } from 'aws-cdk-lib/aws-lambda';
+import * as lambda from 'aws-cdk-lib/aws-lambda-nodejs';
+import { NagSuppressions } from 'cdk-nag';
+import { Construct } from 'constructs';
+
+/** Pickup Lambda timeout (minutes). */
+const PICKUP_TIMEOUT_MINUTES = 5;
+
+/** Pickup Lambda memory (MB). */
+const PICKUP_MEMORY_MB = 256;
+
+/**
+ * Default pickup schedule interval (minutes). Short — queue latency is
+ * user-visible wait time; 1 minute keeps the worst-case pickup delay
+ * bounded while a QUEUED-empty cycle is a single cheap GSI query.
+ */
+const DEFAULT_SCHEDULE_MINUTES = 1;
+
+/** Default max queue age before the backstop fails the task (seconds; 24h). */
+const DEFAULT_QUEUE_MAX_AGE_SECONDS = 86400;
+
+/** Default task-record retention used for event TTL (days). */
+const DEFAULT_TASK_RETENTION_DAYS = 90;
+
+/**
+ * Properties for AdmissionQueuePickup construct.
+ */
+export interface AdmissionQueuePickupProps {
+  /** TaskTable (StatusIndex GSI powers the QUEUED FIFO query). */
+  readonly taskTable: dynamodb.ITable;
+
+  /** TaskEventsTable (handler writes queue_pickup / task_failed events). */
+  readonly taskEventsTable: dynamodb.ITable;
+
+  /** UserConcurrencyTable (read-only capacity check per user). */
+  readonly userConcurrencyTable: dynamodb.ITable;
+
+  /** ARN of the orchestrator Lambda alias to re-invoke on pickup. */
+  readonly orchestratorFunctionArn: string;
+
+  /**
+   * Maximum concurrent tasks per user — must match the orchestrator's
+   * value so the capacity pre-check agrees with `admissionControl`.
+   * @default 10
+   */
+  readonly maxConcurrentTasksPerUser?: number;
+
+  /**
+   * How often to drain the queue.
+   * @default Duration.minutes(1)
+   */
+  readonly schedule?: Duration;
+
+  /**
+   * Max time a task may sit QUEUED before the backstop fails it (seconds).
+   * @default 86400 (24 hours)
+   */
+  readonly queueMaxAgeSeconds?: number;
+
+  /** Forwarded to the handler for event TTL. @default 90 */
+  readonly taskRetentionDays?: number;
+}
+
+/**
+ * Scheduled Lambda that drains the admission queue (#441).
+ *
+ * Tasks that hit the per-user concurrency cap are parked in QUEUED by the
+ * orchestrator instead of FAILED. This Lambda re-attempts admission in
+ * FIFO order (StatusIndex GSI, ascending ``created_at``) as slots free
+ * up: it flips QUEUED -> SUBMITTED and re-invokes the orchestrator, whose
+ * atomic `admissionControl` remains the single writer of the concurrency
+ * counter (a lost race simply re-queues the task, preserving position).
+ */
+export class AdmissionQueuePickup extends Construct {
+  public readonly fn: lambda.NodejsFunction;
+
+  constructor(scope: Construct, id: string, props: AdmissionQueuePickupProps) {
+    super(scope, id);
+
+    const handlersDir = path.join(__dirname, '..', 'handlers');
+
+    this.fn = new lambda.NodejsFunction(this, 'PickupFn', {
+      entry: path.join(handlersDir, 'reconcile-admission-queue.ts'),
+      handler: 'handler',
+      runtime: Runtime.NODEJS_24_X,
+      architecture: Architecture.ARM_64,
+      timeout: Duration.minutes(PICKUP_TIMEOUT_MINUTES),
+      memorySize: PICKUP_MEMORY_MB,
+      environment: {
+        TASK_TABLE_NAME: props.taskTable.tableName,
+        TASK_EVENTS_TABLE_NAME: props.taskEventsTable.tableName,
+        USER_CONCURRENCY_TABLE_NAME: props.userConcurrencyTable.tableName,
+        ORCHESTRATOR_FUNCTION_ARN: props.orchestratorFunctionArn,
+        MAX_CONCURRENT_TASKS_PER_USER: String(props.maxConcurrentTasksPerUser ?? 10),
+        QUEUE_MAX_AGE_SECONDS: String(props.queueMaxAgeSeconds ?? DEFAULT_QUEUE_MAX_AGE_SECONDS),
+        TASK_RETENTION_DAYS: String(props.taskRetentionDays ?? DEFAULT_TASK_RETENTION_DAYS),
+      },
+      bundling: {
+        externalModules: ['@aws-sdk/*'],
+      },
+    });
+
+    // TaskTable: StatusIndex query + conditional QUEUED->SUBMITTED /
+    // QUEUED->FAILED transitions.
+    props.taskTable.grantReadWriteData(this.fn);
+    // TaskEvents: queue_pickup / task_failed events.
+    props.taskEventsTable.grantWriteData(this.fn);
+    // Concurrency: READ-ONLY capacity pre-check — the orchestrator's
+    // admissionControl is the only writer of the counter.
+    props.userConcurrencyTable.grantReadData(this.fn);
+
+    // Re-invoke the orchestrator alias on pickup.
+    this.fn.addToRolePolicy(new iam.PolicyStatement({
+      actions: ['lambda:InvokeFunction'],
+      resources: [props.orchestratorFunctionArn],
+    }));
+
+    const schedule = props.schedule ?? Duration.minutes(DEFAULT_SCHEDULE_MINUTES);
+    const rule = new events.Rule(this, 'PickupSchedule', {
+      schedule: events.Schedule.rate(schedule),
+    });
+    rule.addTarget(new targets.LambdaFunction(this.fn));
+
+    NagSuppressions.addResourceSuppressions(this.fn, [
+      {
+        id: 'AwsSolutions-IAM4',
+        reason: 'AWSLambdaBasicExecutionRole is required for CloudWatch Logs access',
+      },
+      {
+        id: 'AwsSolutions-IAM5',
+        reason: 'DynamoDB index/* wildcards generated by CDK grantReadWriteData/grantReadData for StatusIndex query access',
+      },
+    ], true);
+  }
+}

cdk/src/constructs/task-api.ts

@@ -121,6 +121,15 @@ export interface TaskApiProps {
    */
   readonly orchestratorFunctionArn?: string;
 
+  /**
+   * Maximum concurrent tasks per user (#441). Threaded to the get-task
+   * handler so the queue-wait ETA heuristic agrees with the
+   * orchestrator's admission cap. Must match
+   * ``TaskOrchestrator.maxConcurrentTasksPerUser``.
+   * @default 10
+   */
+  readonly maxConcurrentTasksPerUser?: number;
+
   /**
    * API Gateway stage name.
    * @default 'v1'
@@ -533,7 +542,12 @@ export class TaskApi extends Construct {
       handler: 'handler',
       runtime: Runtime.NODEJS_24_X,
       architecture: Architecture.ARM_64,
-      environment: commonEnv,
+      environment: {
+        ...commonEnv,
+        // #441: queue-position ETA heuristic must agree with the
+        // orchestrator's per-user admission cap.
+        MAX_CONCURRENT_TASKS_PER_USER: String(props.maxConcurrentTasksPerUser ?? 10),
+      },
       bundling: commonBundling,
     });
 

cdk/src/constructs/task-status.ts

@@ -21,14 +21,23 @@
  * Valid task states in the task lifecycle state machine.
  *
  * States progress through the lifecycle:
- *   [PENDING_UPLOADS ->] SUBMITTED -> HYDRATING ->
+ *   [PENDING_UPLOADS ->] SUBMITTED [<-> QUEUED] -> HYDRATING ->
  *   RUNNING -> FINALIZING -> terminal (COMPLETED / FAILED / CANCELLED / TIMED_OUT).
  * See ORCHESTRATOR.md for the full state transition table.
  *
  * PENDING_UPLOADS is a pre-active state for tasks with presigned-upload
  * attachments: no compute allocated, no concurrency slot consumed. The
  * task transitions to SUBMITTED once uploads are confirmed and screened.
  *
+ * QUEUED (#441) is a pre-active state for tasks that hit the per-user
+ * admission cap: the admission slot was NOT acquired, so no compute is
+ * allocated and no concurrency slot is consumed. The admission-queue
+ * pickup Lambda re-attempts admission in FIFO order (by ``created_at``)
+ * as slots free up, transitioning QUEUED -> SUBMITTED and re-invoking
+ * the orchestrator. A pickup that loses the admission race simply
+ * re-queues (SUBMITTED -> QUEUED); FIFO position is preserved because
+ * ``created_at`` never changes.
+ *
  * AWAITING_APPROVAL is the Cedar-HITL soft-deny gate surface: the
  * task is alive but paused on a human decision. See
  * `docs/design/CEDAR_HITL_GATES.md` §10.3 for the joint
@@ -37,6 +46,7 @@
  */
 export const TaskStatus = {
   PENDING_UPLOADS: 'PENDING_UPLOADS',
+  QUEUED: 'QUEUED',
   SUBMITTED: 'SUBMITTED',
   HYDRATING: 'HYDRATING',
   RUNNING: 'RUNNING',
@@ -66,10 +76,14 @@ export const TERMINAL_STATUSES: readonly TaskStatusType[] = [
 /**
  * Pre-active states where the task exists but has not entered the
  * orchestration pipeline. No compute resources are allocated and no
- * concurrency slot is consumed.
+ * concurrency slot is consumed. QUEUED belongs here (#441): admission
+ * explicitly did NOT acquire a slot, so counting it as active would
+ * deadlock the queue (queued tasks would hold the very slots they
+ * are waiting for).
  */
 export const PRE_ACTIVE_STATUSES: readonly TaskStatusType[] = [
   TaskStatus.PENDING_UPLOADS,
+  TaskStatus.QUEUED,
 ];
 
 /**
@@ -97,7 +111,11 @@ export const VALID_TRANSITIONS: Readonly<Record<TaskStatusType, readonly TaskSta
   // Transitions to SUBMITTED on confirm-uploads success, FAILED on
   // screening failure, CANCELLED on user cancel or 30-min auto-cancel.
   [TaskStatus.PENDING_UPLOADS]: [TaskStatus.SUBMITTED, TaskStatus.FAILED, TaskStatus.CANCELLED],
-  [TaskStatus.SUBMITTED]: [TaskStatus.HYDRATING, TaskStatus.FAILED, TaskStatus.CANCELLED],
+  // QUEUED (#441): admission-capped task awaiting a free concurrency
+  // slot. SUBMITTED on pickup (slot acquired), CANCELLED on user
+  // cancel, FAILED only via the queue-stranded backstop.
+  [TaskStatus.QUEUED]: [TaskStatus.SUBMITTED, TaskStatus.CANCELLED, TaskStatus.FAILED],
+  [TaskStatus.SUBMITTED]: [TaskStatus.QUEUED, TaskStatus.HYDRATING, TaskStatus.FAILED, TaskStatus.CANCELLED],
   [TaskStatus.HYDRATING]: [
     TaskStatus.RUNNING,
     TaskStatus.AWAITING_APPROVAL,

cdk/src/handlers/get-task.ts

@@ -18,9 +18,10 @@
  */
 
 import { DynamoDBClient } from '@aws-sdk/client-dynamodb';
-import { DynamoDBDocumentClient, GetCommand } from '@aws-sdk/lib-dynamodb';
+import { DynamoDBDocumentClient, GetCommand, QueryCommand } from '@aws-sdk/lib-dynamodb';
 import type { APIGatewayProxyEvent, APIGatewayProxyResult } from 'aws-lambda';
 import { ulid } from 'ulid';
+import { TaskStatus } from '../constructs/task-status';
 import { extractUserId } from './shared/gateway';
 import { logger } from './shared/logger';
 import { ErrorCode, errorResponse, successResponse } from './shared/response';
@@ -29,6 +30,81 @@ import { type TaskRecord, toTaskDetail } from './shared/types';
 const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({}));
 const TABLE_NAME = process.env.TASK_TABLE_NAME!;
 
+/** Must match the orchestrator's per-user admission cap for a sane ETA. */
+const MAX_CONCURRENT = Number(process.env.MAX_CONCURRENT_TASKS_PER_USER ?? '10');
+
+/**
+ * Rough per-task duration used for the queue-wait heuristic (#441).
+ * Deliberately coarse — the ETA is a UX hint ("minutes, not seconds"),
+ * not a promise. Override via env for fleets with unusual task profiles.
+ */
+const QUEUE_ETA_AVG_TASK_DURATION_S = Number(process.env.QUEUE_ETA_AVG_TASK_DURATION_S ?? '600');
+
+/**
+ * Compute the task's 1-based FIFO position among the user's QUEUED tasks
+ * (#441) plus a coarse wait estimate. Position is per user because the
+ * admission cap is per user — a global position would overstate the wait
+ * for users whose slots are free.
+ *
+ * Queries the UserStatusIndex GSI (PK user_id, SK status_created_at
+ * ``QUEUED#...``) and ranks by ``created_at``, which is the pickup
+ * Lambda's FIFO key (``status_created_at`` moves on re-queue;
+ * ``created_at`` never does).
+ *
+ * Fail-open: any error returns undefined so a GSI hiccup degrades the
+ * response to ``queue_position: null`` instead of failing the whole GET.
+ */
+async function computeQueueInfo(
+  record: TaskRecord,
+): Promise<{ queue_position: number; estimated_wait_s: number | null } | undefined> {
+  try {
+    let ahead = 0;
+    let found = false;
+    let lastKey: Record<string, unknown> | undefined;
+    do {
+      const resp = await ddb.send(new QueryCommand({
+        TableName: TABLE_NAME,
+        IndexName: 'UserStatusIndex',
+        KeyConditionExpression: 'user_id = :uid AND begins_with(status_created_at, :queued)',
+        ExpressionAttributeValues: {
+          ':uid': record.user_id,
+          ':queued': `${TaskStatus.QUEUED}#`,
+        },
+        ProjectionExpression: 'task_id, created_at',
+        ExclusiveStartKey: lastKey as Record<string, never> | undefined,
+      }));
+      for (const item of resp.Items ?? []) {
+        if (item.task_id === record.task_id) {
+          found = true;
+        } else if (typeof item.created_at === 'string' && item.created_at < record.created_at) {
+          ahead++;
+        }
+      }
+      lastKey = resp.LastEvaluatedKey;
+    } while (lastKey);
+
+    // The task left QUEUED between our GetItem and this query — report
+    // no queue info rather than a stale position.
+    if (!found) {
+      return undefined;
+    }
+
+    const position = ahead + 1;
+    // Coarse ETA: slots drain MAX_CONCURRENT at a time, each batch
+    // lasting roughly one average task duration.
+    const estimatedWaitS = MAX_CONCURRENT > 0 && QUEUE_ETA_AVG_TASK_DURATION_S > 0
+      ? Math.ceil(position / MAX_CONCURRENT) * QUEUE_ETA_AVG_TASK_DURATION_S
+      : null;
+    return { queue_position: position, estimated_wait_s: estimatedWaitS };
+  } catch (err) {
+    logger.warn('Queue position lookup failed (fail-open — returning null position)', {
+      task_id: record.task_id,
+      error: err instanceof Error ? err.message : String(err),
+    });
+    return undefined; // nosemgrep: ts-silent-success-masking -- queue position is a best-effort UX hint; a GSI failure must not fail the whole GET /tasks/{id}
+  }
+}
+
 /**
  * GET /v1/tasks/{task_id} — Get full task details.
  */
@@ -64,8 +140,13 @@ export async function handler(event: APIGatewayProxyEvent): Promise<APIGatewayPr
       return errorResponse(403, ErrorCode.FORBIDDEN, 'You do not have access to this task.', requestId);
     }
 
-    // 5. Return task detail
-    return successResponse(200, toTaskDetail(record), requestId);
+    // 5. For QUEUED tasks, compute read-time queue position + ETA (#441)
+    const queueInfo = record.status === TaskStatus.QUEUED
+      ? await computeQueueInfo(record)
+      : undefined;
+
+    // 6. Return task detail
+    return successResponse(200, toTaskDetail(record, queueInfo), requestId);
   } catch (err) {
     logger.error('Failed to get task', { error: String(err), request_id: requestId });
     return errorResponse(500, ErrorCode.INTERNAL_ERROR, 'Internal server error.', requestId);