From 74d3f5768beae71d55227f5627979320525d924a Mon Sep 17 00:00:00 2001 From: bgagent Date: Tue, 7 Jul 2026 16:23:20 -0400 Subject: [PATCH 1/2] feat(orchestration): admission queue with deferred pickup (#441) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When a task hits the per-user concurrency cap, park it in a new QUEUED state instead of failing it (the #331 mass-fail mode). A scheduled AdmissionQueuePickup Lambda drains the queue FIFO (by created_at) as slots free up, flipping QUEUED -> SUBMITTED and re-invoking the orchestrator, whose atomic admissionControl stays the single writer of the concurrency counter — a lost pickup race harmlessly re-queues without losing FIFO position. - task-status: new QUEUED state (pre-active — holds no slot), with SUBMITTED <-> QUEUED transitions plus QUEUED -> CANCELLED/FAILED - orchestrator: queueTask() replaces failTask on cap; queued_at is stamped once, admission_attempts increments per pass; Linear/Jira channel feedback now says 'queued' and fires only on first entry - reconcile-admission-queue: FIFO drain per user with read-only capacity pre-check, conditional flips (cancel-safe), 24h max-age backstop, and orchestrator re-invoke carrying a pickup nonce - reconcile-stranded-tasks: age by time-in-current-status so a task that waited in the queue longer than the stranded timeout is not killed right after pickup - API/CLI: GET /tasks/{id} computes read-time queue_position + estimated_wait_s (fail-open); bgagent status/detail render a 'Queue: position N (est. wait ~Xm)' line - cancel: QUEUED is cancellable everywhere (REST already generic; Slack cancel list extended); no concurrency release needed - idempotent replay returns the existing QUEUED task unchanged - tests: state-machine invariants, queueTask, pickup handler (incl. #331 fan-out-burst regression: burst above cap survives and drains FIFO with zero failures), get-task queue position, CLI rendering Closes #441 --- cdk/src/constructs/admission-queue-pickup.ts | 160 +++++++ cdk/src/constructs/task-api.ts | 16 +- cdk/src/constructs/task-status.ts | 24 +- cdk/src/handlers/get-task.ts | 87 +++- cdk/src/handlers/orchestrate-task.ts | 75 ++-- cdk/src/handlers/reconcile-admission-queue.ts | 384 +++++++++++++++++ cdk/src/handlers/reconcile-stranded-tasks.ts | 18 +- cdk/src/handlers/shared/orchestrator.ts | 47 ++ cdk/src/handlers/shared/types.ts | 36 +- cdk/src/handlers/slack-interactions.ts | 8 +- cdk/src/stacks/agent.ts | 13 + .../constructs/admission-queue-pickup.test.ts | 120 ++++++ cdk/test/constructs/task-status.test.ts | 54 ++- cdk/test/handlers/get-task.test.ts | 75 ++++ cdk/test/handlers/orchestrate-task.test.ts | 56 +++ .../reconcile-admission-queue.test.ts | 404 ++++++++++++++++++ cli/src/format.ts | 19 + cli/src/types.ts | 13 + cli/test/format-status-snapshot.test.ts | 22 + cli/test/format.test.ts | 36 ++ docs/design/ORCHESTRATOR.md | 19 +- .../content/docs/architecture/Orchestrator.md | 19 +- 22 files changed, 1648 insertions(+), 57 deletions(-) create mode 100644 cdk/src/constructs/admission-queue-pickup.ts create mode 100644 cdk/src/handlers/reconcile-admission-queue.ts create mode 100644 cdk/test/constructs/admission-queue-pickup.test.ts create mode 100644 cdk/test/handlers/reconcile-admission-queue.test.ts diff --git a/cdk/src/constructs/admission-queue-pickup.ts b/cdk/src/constructs/admission-queue-pickup.ts new file mode 100644 index 00000000..162235a1 --- /dev/null +++ b/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); + } +} diff --git a/cdk/src/constructs/task-api.ts b/cdk/src/constructs/task-api.ts index d2f67993..d10d3ecb 100644 --- a/cdk/src/constructs/task-api.ts +++ b/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, }); diff --git a/cdk/src/constructs/task-status.ts b/cdk/src/constructs/task-status.ts index a5db1e33..3451155c 100644 --- a/cdk/src/constructs/task-status.ts +++ b/cdk/src/constructs/task-status.ts @@ -21,7 +21,7 @@ * 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. * @@ -29,6 +29,15 @@ * 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 { + try { + let ahead = 0; + let found = false; + let lastKey: Record | 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 | 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 SUBMITTED flip. Unused by the + * handler; it exists so the re-invoke payload differs from the + * original create-task invocation and no layer (Lambda async dedup, + * durable-execution idempotency) can mistake it for a replay. + */ + readonly queue_pickup_id?: string; } const MAX_POLL_ATTEMPTS = 1020; // ~8.5h at 30s intervals @@ -68,7 +77,10 @@ const durableHandler: DurableExecutionHandler = asyn } }); - // Step 2: Admission control — check concurrency limit + // Step 2: Admission control — check concurrency limit. At capacity the + // task is QUEUED (durable, FIFO on created_at) instead of FAILED (#441): + // the admission-queue pickup Lambda re-attempts admission as slots free + // up, flips QUEUED -> SUBMITTED, and re-invokes this orchestrator. const admitted = await context.step('admission-control', async () => { // Re-read status to detect external cancellation between steps const current = await loadTask(taskId); @@ -77,25 +89,31 @@ const durableHandler: DurableExecutionHandler = asyn } const result = await admissionControl(task); if (!result) { - await failTask(taskId, current.status, 'User concurrency limit reached', task.user_id, false); - await emitTaskEvent(taskId, 'admission_rejected', { reason: 'concurrency_limit' }); - // Channel feedback is non-fatal: a throw here would re-run failTask + - // emitTaskEvent on the durable-execution retry, producing duplicate events. - try { - await notifyLinearOnConcurrencyCap(task); - } catch (err) { - logger.warn('Linear concurrency-cap feedback failed (non-fatal)', { - task_id: taskId, - error: err instanceof Error ? err.message : String(err), - }); - } - try { - await notifyJiraOnConcurrencyCap(task); - } catch (err) { - logger.warn('Jira concurrency-cap feedback failed (non-fatal)', { - task_id: taskId, - error: err instanceof Error ? err.message : String(err), - }); + // Re-read for queue bookkeeping: `task` is the step-1 snapshot, but a + // pickup-cycle re-queue must preserve the CURRENT queued_at / + // admission_attempts, which the pickup Lambda may have advanced. + await queueTask(current); + // Channel feedback only on the FIRST queue entry — a pickup-cycle + // re-queue (lost admission race) would otherwise spam the issue with + // one comment per cycle. Non-fatal: a throw here would re-run + // queueTask on the durable-execution retry, producing duplicate events. + if (current.queued_at === undefined) { + try { + await notifyLinearOnConcurrencyCap(task); + } catch (err) { + logger.warn('Linear concurrency-cap feedback failed (non-fatal)', { + task_id: taskId, + error: err instanceof Error ? err.message : String(err), + }); + } + try { + await notifyJiraOnConcurrencyCap(task); + } catch (err) { + logger.warn('Jira concurrency-cap feedback failed (non-fatal)', { + task_id: taskId, + error: err instanceof Error ? err.message : String(err), + }); + } } } return result; @@ -303,13 +321,12 @@ const durableHandler: DurableExecutionHandler = asyn export const handler = withDurableExecution(durableHandler); /** - * Post a Linear comment + ❌ reaction when admission control rejects a task - * for the user concurrency cap. Linear-only; silently no-ops for other - * channels. + * Post a Linear comment when admission control queues a task on the user + * concurrency cap (#441). Linear-only; silently no-ops for other channels. * * The processor side (`linear-webhook-processor.ts`) already covers * pre-`createTaskCore` rejections (unmapped project, unlinked actor, guardrail); - * this hook covers the post-201 case where the orchestrator rejects on + * this hook covers the post-201 case where the orchestrator queues on * admission. Without this, the only Linear-side signal would be the 👀 * reaction the agent never gets to add — looks like the integration silently * dropped the request. @@ -342,7 +359,7 @@ export async function notifyLinearOnConcurrencyCap(task: TaskRecord): Promise QUEUED), + * `created_at` never changes, so FIFO position is preserved for the next + * cycle. This at-most-bounces-once-per-cycle design avoids a second + * writer on the counter (the #331 class of drift bugs). + * + * Cancel semantics: the QUEUED -> SUBMITTED flip is conditional on the + * task still being QUEUED, so a user cancel between the query and the + * flip cleanly wins and the task is skipped. + * + * Backstop: a task QUEUED longer than ``QUEUE_MAX_AGE_SECONDS`` is failed + * (QUEUED -> FAILED) with an explanatory message so the queue can never + * accumulate unbounded zombies (e.g. a user's slots wedged by a + * counter-drift bug that the concurrency reconciler hasn't corrected yet). + * No concurrency slot is released — QUEUED tasks never held one. + */ + +import { + DynamoDBClient, + GetItemCommand, + PutItemCommand, + QueryCommand, + UpdateItemCommand, +} from '@aws-sdk/client-dynamodb'; +import { InvokeCommand, LambdaClient } from '@aws-sdk/client-lambda'; +import { ulid } from 'ulid'; +import { logger } from './shared/logger'; + +const ddb = new DynamoDBClient({}); +const lambdaClient = new LambdaClient({}); + +const TASK_TABLE = process.env.TASK_TABLE_NAME!; +const EVENTS_TABLE = process.env.TASK_EVENTS_TABLE_NAME!; +const CONCURRENCY_TABLE = process.env.USER_CONCURRENCY_TABLE_NAME!; +const ORCHESTRATOR_FUNCTION_ARN = process.env.ORCHESTRATOR_FUNCTION_ARN!; +const MAX_CONCURRENT = Number(process.env.MAX_CONCURRENT_TASKS_PER_USER ?? '10'); +const TASK_RETENTION_DAYS = Number(process.env.TASK_RETENTION_DAYS ?? '90'); + +/** + * Maximum time a task may sit QUEUED before the backstop fails it. + * Default 24h — far above any legitimate queue wait (tasks run minutes + * to a few hours), so tripping it indicates wedged capacity, not load. + */ +const QUEUE_MAX_AGE_SECONDS = Number(process.env.QUEUE_MAX_AGE_SECONDS ?? '86400'); + +interface QueuedTask { + readonly task_id: string; + readonly user_id: string; + readonly created_at: string; + readonly age_seconds: number; + readonly admission_attempts: number; +} + +/** Result counters for the final log line / test assertions. */ +export interface PickupSummary { + queued_seen: number; + picked_up: number; + expired: number; + skipped_no_capacity: number; + skipped_race: number; + errors: number; +} + +/** + * Query ALL QUEUED tasks in global FIFO order (StatusIndex GSI: PK + * ``status``, SK ``created_at`` — ascending scan gives oldest-first). + */ +async function listQueuedTasks(now: Date): Promise { + const tasks: QueuedTask[] = []; + let lastKey: Record | undefined; + + do { + const resp = await ddb.send(new QueryCommand({ + TableName: TASK_TABLE, + IndexName: 'StatusIndex', + KeyConditionExpression: '#s = :queued', + ExpressionAttributeNames: { '#s': 'status' }, + ExpressionAttributeValues: { ':queued': { S: 'QUEUED' } }, + // ScanIndexForward defaults to true — ascending created_at, i.e. FIFO. + ExclusiveStartKey: lastKey as Record | undefined, + })); + + for (const item of resp.Items ?? []) { + const taskId = item.task_id?.S; + const userId = item.user_id?.S; + const createdAt = item.created_at?.S; + if (!taskId || !userId || !createdAt) continue; + const createdMs = Date.parse(createdAt); + tasks.push({ + task_id: taskId, + user_id: userId, + created_at: createdAt, + age_seconds: Number.isNaN(createdMs) ? 0 : Math.floor((now.getTime() - createdMs) / 1000), + admission_attempts: Number(item.admission_attempts?.N ?? '0'), + }); + } + + lastKey = resp.LastEvaluatedKey; + } while (lastKey); + + return tasks; +} + +/** Read a user's current active_count (0 when the row does not exist). */ +async function readActiveCount(userId: string): Promise { + const resp = await ddb.send(new GetItemCommand({ + TableName: CONCURRENCY_TABLE, + Key: { user_id: { S: userId } }, + ProjectionExpression: 'active_count', + })); + return Number(resp.Item?.active_count?.N ?? '0'); +} + +/** Best-effort TaskEvents write — event loss is acceptable, the task record is the source of truth. */ +async function emitEvent(taskId: string, eventType: string, metadata: Record): Promise { + const ttl = Math.floor(Date.now() / 1000) + TASK_RETENTION_DAYS * 24 * 3600; + try { + await ddb.send(new PutItemCommand({ + TableName: EVENTS_TABLE, + Item: { + task_id: { S: taskId }, + event_id: { S: ulid() }, + event_type: { S: eventType }, + timestamp: { S: new Date().toISOString() }, + ttl: { N: String(ttl) }, + metadata: { M: metadata }, + }, + })); + } catch (err) { + logger.warn(`Failed to write ${eventType} event (best-effort)`, { + task_id: taskId, + error: err instanceof Error ? err.message : String(err), + }); + } +} + +/** + * Flip one task QUEUED -> SUBMITTED and re-invoke the orchestrator. + * Returns false when the conditional flip lost to a concurrent + * transition (user cancel / another pickup instance). + */ +async function pickUpTask(task: QueuedTask): Promise { + const now = new Date().toISOString(); + try { + await ddb.send(new UpdateItemCommand({ + TableName: TASK_TABLE, + Key: { task_id: { S: task.task_id } }, + UpdateExpression: 'SET #s = :submitted, updated_at = :now, status_created_at = :sca', + ConditionExpression: '#s = :queued', + ExpressionAttributeNames: { '#s': 'status' }, + ExpressionAttributeValues: { + ':submitted': { S: 'SUBMITTED' }, + ':queued': { S: 'QUEUED' }, + ':now': { S: now }, + ':sca': { S: `SUBMITTED#${now}` }, + }, + })); + } catch (err: unknown) { + if (err && typeof err === 'object' && 'name' in err && err.name === 'ConditionalCheckFailedException') { + logger.info('Queued task transitioned concurrently before pickup — skipping', { task_id: task.task_id }); + return false; + } + throw err; + } + + await emitEvent(task.task_id, 'queue_pickup', { + queued_for_s: { N: String(task.age_seconds) }, + admission_attempts: { N: String(task.admission_attempts) }, + }); + + // Async re-invoke. The payload carries a per-pickup nonce so no layer + // (Lambda async dedup, durable-execution idempotency) can mistake the + // re-invoke for a replay of the original create-task invocation. + try { + await lambdaClient.send(new InvokeCommand({ + FunctionName: ORCHESTRATOR_FUNCTION_ARN, + InvocationType: 'Event', + Payload: new TextEncoder().encode(JSON.stringify({ + task_id: task.task_id, + queue_pickup_id: ulid(), + })), + })); + } catch (invokeErr) { + // The task is now SUBMITTED with no pipeline attached. Do NOT try to + // roll back (racy) — the stranded-task reconciler sweeps SUBMITTED + // tasks with no pipeline, which is exactly this failure mode. + logger.error('Orchestrator re-invoke failed after queue pickup — stranded-task reconciler will sweep', { + task_id: task.task_id, + error: invokeErr instanceof Error ? invokeErr.message : String(invokeErr), + }); + return false; + } + + logger.info('Queued task picked up', { + task_id: task.task_id, + user_id: task.user_id, + queued_for_s: task.age_seconds, + }); + return true; +} + +/** + * Backstop: fail a task that has been QUEUED past the max age. No + * concurrency release — QUEUED tasks never hold a slot. + */ +async function expireQueuedTask(task: QueuedTask): Promise { + const now = new Date().toISOString(); + const errorMessage = + `Admission queue timeout: task waited ${task.age_seconds}s for a free ` + + 'concurrency slot (limit per user) without being admitted. This usually ' + + 'means long-running tasks are holding all slots — cancel one or wait, ' + + 'then resubmit.'; + const ttl = Math.floor(Date.now() / 1000) + TASK_RETENTION_DAYS * 24 * 3600; + try { + await ddb.send(new UpdateItemCommand({ + TableName: TASK_TABLE, + Key: { task_id: { S: task.task_id } }, + UpdateExpression: + 'SET #s = :failed, updated_at = :now, completed_at = :now, ' + + 'error_message = :err, status_created_at = :sca, #ttl = :ttl', + ConditionExpression: '#s = :queued', + ExpressionAttributeNames: { '#s': 'status', '#ttl': 'ttl' }, + ExpressionAttributeValues: { + ':failed': { S: 'FAILED' }, + ':queued': { S: 'QUEUED' }, + ':now': { S: now }, + ':err': { S: errorMessage }, + ':sca': { S: `FAILED#${now}` }, + ':ttl': { N: String(ttl) }, + }, + })); + } catch (err: unknown) { + if (err && typeof err === 'object' && 'name' in err && err.name === 'ConditionalCheckFailedException') { + return false; + } + throw err; + } + await emitEvent(task.task_id, 'task_failed', { + error_message: { S: errorMessage }, + reason: { S: 'queue_timeout' }, + }); + logger.warn('Queued task expired by backstop', { + task_id: task.task_id, + user_id: task.user_id, + age_seconds: task.age_seconds, + }); + return true; +} + +export async function handler(): Promise { + const now = new Date(); + const summary: PickupSummary = { + queued_seen: 0, + picked_up: 0, + expired: 0, + skipped_no_capacity: 0, + skipped_race: 0, + errors: 0, + }; + + let queued: QueuedTask[]; + try { + queued = await listQueuedTasks(now); + } catch (err) { + logger.error('Admission-queue query failed — aborting cycle', { + error: err instanceof Error ? err.message : String(err), + }); + throw err; + } + summary.queued_seen = queued.length; + if (queued.length === 0) { + return summary; + } + + // Group per user, preserving the global FIFO order within each group. + const byUser = new Map(); + for (const task of queued) { + const list = byUser.get(task.user_id); + if (list) { + list.push(task); + } else { + byUser.set(task.user_id, [task]); + } + } + + for (const [userId, userTasks] of byUser) { + // Expire over-age tasks first so they never block capacity math. + const live: QueuedTask[] = []; + for (const task of userTasks) { + if (task.age_seconds > QUEUE_MAX_AGE_SECONDS) { + try { + if (await expireQueuedTask(task)) { + summary.expired++; + } else { + summary.skipped_race++; + } + } catch (err) { + summary.errors++; + logger.warn('Failed to expire over-age queued task, continuing', { + task_id: task.task_id, + error: err instanceof Error ? err.message : String(err), + }); + } + } else { + live.push(task); + } + } + if (live.length === 0) continue; + + let available: number; + try { + available = MAX_CONCURRENT - await readActiveCount(userId); + } catch (err) { + summary.errors++; + logger.warn('Failed to read concurrency for user — skipping their queue this cycle', { + user_id: userId, + error: err instanceof Error ? err.message : String(err), + }); + continue; + } + + if (available <= 0) { + summary.skipped_no_capacity += live.length; + continue; + } + + // Oldest-first pickup, bounded by free capacity. The orchestrator's + // atomic admissionControl is still the final arbiter — an optimistic + // over-pick simply re-queues. + for (const task of live.slice(0, available)) { + try { + if (await pickUpTask(task)) { + summary.picked_up++; + } else { + summary.skipped_race++; + } + } catch (err) { + summary.errors++; + logger.warn('Per-task queue pickup failed, continuing', { + task_id: task.task_id, + error: err instanceof Error ? err.message : String(err), + }); + } + } + summary.skipped_no_capacity += Math.max(0, live.length - available); + } + + const level = summary.errors > 0 && summary.picked_up === 0 && summary.queued_seen > 0 ? 'error' : 'info'; + logger[level]('Admission-queue pickup finished', { ...summary }); + return summary; +} diff --git a/cdk/src/handlers/reconcile-stranded-tasks.ts b/cdk/src/handlers/reconcile-stranded-tasks.ts index 8688601f..73f826b4 100644 --- a/cdk/src/handlers/reconcile-stranded-tasks.ts +++ b/cdk/src/handlers/reconcile-stranded-tasks.ts @@ -123,8 +123,22 @@ async function findStrandedCandidates( const createdAt = item.created_at?.S; if (!taskId || !userId || !createdAt) continue; - const createdMs = Date.parse(createdAt); - const ageSec = Math.floor((now.getTime() - createdMs) / 1000); + // Age by time-in-CURRENT-status, not creation time (#441). A task + // that waited in the admission queue longer than the stranded + // timeout and was then picked up (QUEUED -> SUBMITTED) has an old + // created_at but a fresh status entry — failing it here would kill + // it before its pipeline attaches. status_created_at is + // `#`; created_at <= status-entry time always, so the + // GSI cutoff on created_at remains a correct superset pre-filter. + // Records missing/with a malformed status_created_at fall back to + // created_at (pre-#441 behavior). + const statusEnteredAt = item.status_created_at?.S?.split('#')[1]; + const ageAnchor = statusEnteredAt && !Number.isNaN(Date.parse(statusEnteredAt)) + ? statusEnteredAt + : createdAt; + const anchorMs = Date.parse(ageAnchor); + const ageSec = Math.floor((now.getTime() - anchorMs) / 1000); + if (ageSec < timeoutSeconds) continue; matches.push({ task_id: taskId, diff --git a/cdk/src/handlers/shared/orchestrator.ts b/cdk/src/handlers/shared/orchestrator.ts index 27a8a0a6..2274d89c 100644 --- a/cdk/src/handlers/shared/orchestrator.ts +++ b/cdk/src/handlers/shared/orchestrator.ts @@ -832,6 +832,53 @@ export async function finalizeTask( await decrementConcurrency(userId); } +/** + * Queue a task that hit the per-user admission cap (#441). Transitions + * SUBMITTED -> QUEUED and stamps queue bookkeeping; the admission-queue + * pickup Lambda later re-attempts admission in FIFO order (by + * ``created_at``) as slots free up. + * + * ``queued_at`` is written only on the FIRST queue entry — a re-queue + * after a lost admission race preserves the original timestamp. + * ``admission_attempts`` increments on every pass through admission so + * repeatedly-losing tasks are visible to operators. + * + * No concurrency slot is held while QUEUED (admission explicitly did + * not acquire one), so there is nothing to release here. + * + * @param task - the task record (status must be SUBMITTED). + * @returns true if the task was queued; false if a concurrent + * transition (e.g. user cancel) won the conditional check. + */ +export async function queueTask(task: TaskRecord): Promise { + try { + await transitionTask(task.task_id, TaskStatus.SUBMITTED, TaskStatus.QUEUED, { + ...(task.queued_at === undefined && { queued_at: new Date().toISOString() }), + admission_attempts: (typeof task.admission_attempts === 'number' ? task.admission_attempts : 0) + 1, + }); + } catch (err) { + // A concurrent transition (user cancel between our status read and + // this write) wins; the task is no longer SUBMITTED so it must not + // be queued. Anything else is unexpected — rethrow so the durable + // step retries. + if (err && typeof err === 'object' && 'name' in err && err.name === 'ConditionalCheckFailedException') { + logger.info('Queue transition lost to a concurrent status change — skipping', { task_id: task.task_id }); + return false; + } + throw err; + } + await emitTaskEvent(task.task_id, 'admission_queued', { + reason: 'concurrency_limit', + admission_attempts: (typeof task.admission_attempts === 'number' ? task.admission_attempts : 0) + 1, + }); + logger.info('Task queued on admission cap', { + task_id: task.task_id, + user_id: task.user_id, + prior_attempts: task.admission_attempts ?? 0, + }); + return true; +} + /** * Fail a task and release concurrency. Used when admission or hydration fails. * @param taskId - the task ID. diff --git a/cdk/src/handlers/shared/types.ts b/cdk/src/handlers/shared/types.ts index 35d9a8a3..ec5486fc 100644 --- a/cdk/src/handlers/shared/types.ts +++ b/cdk/src/handlers/shared/types.ts @@ -215,6 +215,19 @@ export interface TaskRecord { * atomically on resume (§10.2, §9). */ readonly awaiting_approval_request_id?: string; + /** + * Admission queue (#441): ISO timestamp of the task's FIRST entry + * into QUEUED (set once; a re-queue after a lost admission race does + * not overwrite it). Used for observability — FIFO order itself is + * keyed on ``created_at``, which never changes. + */ + readonly queued_at?: string; + /** + * Admission queue (#441): number of admission attempts made by the + * queue pickup Lambda. Incremented on each pickup attempt; used to + * spot tasks that repeatedly lose the admission race. + */ + readonly admission_attempts?: number; } /** Per-channel override for one notification channel. See @@ -314,6 +327,18 @@ export interface TaskDetail { /** Cedar HITL: when ``status = AWAITING_APPROVAL``, the * ``request_id`` of the pending approval row. Null otherwise. */ readonly awaiting_approval_request_id: string | null; + /** Admission queue (#441): ISO timestamp the task first entered + * QUEUED; null for tasks that were admitted directly. */ + readonly queued_at: string | null; + /** Admission queue (#441): 1-based FIFO position among the caller's + * QUEUED tasks when ``status = QUEUED``. Computed at read time by + * the get-task handler (never persisted — it changes as the queue + * drains); null otherwise or when the handler could not compute it. */ + readonly queue_position: number | null; + /** Admission queue (#441): rough ETA (seconds) until this task is + * picked up, derived from queue_position × the recent average task + * duration heuristic. Null when queue_position is null. */ + readonly estimated_wait_s: number | null; } /** @@ -681,9 +706,15 @@ export interface AgentAttachmentPayload { * the helper when adding new numeric fields. * * @param record - the DynamoDB task record. + * @param queueInfo - read-time queue position/ETA for QUEUED tasks (#441). + * Computed by the caller (get-task) because position changes as the + * queue drains and must never be persisted; omitted everywhere else. * @returns the API-facing task detail. */ -export function toTaskDetail(record: TaskRecord): TaskDetail { +export function toTaskDetail( + record: TaskRecord, + queueInfo?: { readonly queue_position: number; readonly estimated_wait_s: number | null }, +): TaskDetail { const ctx = { task_id: record.task_id }; return { task_id: record.task_id, @@ -737,6 +768,9 @@ export function toTaskDetail(record: TaskRecord): TaskDetail { logger, ), awaiting_approval_request_id: record.awaiting_approval_request_id ?? null, + queued_at: record.queued_at ?? null, + queue_position: queueInfo?.queue_position ?? null, + estimated_wait_s: queueInfo?.estimated_wait_s ?? null, }; } diff --git a/cdk/src/handlers/slack-interactions.ts b/cdk/src/handlers/slack-interactions.ts index ac25ccff..93a8b9d1 100644 --- a/cdk/src/handlers/slack-interactions.ts +++ b/cdk/src/handlers/slack-interactions.ts @@ -131,14 +131,15 @@ async function handleCancelAction(payload: SlackInteractionPayload, actionId: st return; } - // Attempt to cancel. - const CANCELLABLE_STATUSES = ['PENDING_UPLOADS', 'SUBMITTED', 'HYDRATING', 'RUNNING', 'AWAITING_APPROVAL', 'FINALIZING']; + // Attempt to cancel. QUEUED (#441) is cancellable — it removes the + // task from the admission queue (no compute or concurrency to release). + const CANCELLABLE_STATUSES = ['PENDING_UPLOADS', 'QUEUED', 'SUBMITTED', 'HYDRATING', 'RUNNING', 'AWAITING_APPROVAL', 'FINALIZING']; try { await ddb.send(new UpdateCommand({ TableName: TASK_TABLE, Key: { task_id: taskId }, UpdateExpression: 'SET #s = :cancelled, updated_at = :now', - ConditionExpression: '#s IN (:s1, :s2, :s3, :s4, :s5, :s6)', + ConditionExpression: '#s IN (:s1, :s2, :s3, :s4, :s5, :s6, :s7)', ExpressionAttributeNames: { '#s': 'status' }, ExpressionAttributeValues: { ':cancelled': 'CANCELLED', @@ -149,6 +150,7 @@ async function handleCancelAction(payload: SlackInteractionPayload, actionId: st ':s4': CANCELLABLE_STATUSES[3], ':s5': CANCELLABLE_STATUSES[4], ':s6': CANCELLABLE_STATUSES[5], + ':s7': CANCELLABLE_STATUSES[6], }, })); diff --git a/cdk/src/stacks/agent.ts b/cdk/src/stacks/agent.ts index 4e0612d9..9707ff36 100644 --- a/cdk/src/stacks/agent.ts +++ b/cdk/src/stacks/agent.ts @@ -30,6 +30,7 @@ import * as secretsmanager from 'aws-cdk-lib/aws-secretsmanager'; import * as cr from 'aws-cdk-lib/custom-resources'; import { NagSuppressions } from 'cdk-nag'; import { Construct, IConstruct } from 'constructs'; +import { AdmissionQueuePickup } from '../constructs/admission-queue-pickup'; import { AgentMemory } from '../constructs/agent-memory'; import { AgentSessionRole } from '../constructs/agent-session-role'; import { AgentVpc } from '../constructs/agent-vpc'; @@ -639,6 +640,18 @@ export class AgentStack extends Stack { userConcurrencyTable: userConcurrencyTable.table, }); + // --- Admission-queue pickup (#441) --- + // Drains QUEUED tasks (parked by the orchestrator when the per-user + // concurrency cap is hit) in FIFO order as slots free up: flips + // QUEUED -> SUBMITTED and re-invokes the orchestrator, whose atomic + // admissionControl remains the single writer of the counter. + new AdmissionQueuePickup(this, 'AdmissionQueuePickup', { + taskTable: taskTable.table, + taskEventsTable: taskEventsTable.table, + userConcurrencyTable: userConcurrencyTable.table, + orchestratorFunctionArn: orchestrator.alias.functionArn, + }); + // --- Stranded-task reconciler --- // Catches SUBMITTED / HYDRATING tasks whose pipeline never started // (orchestrator Lambda crash between TaskTable write and InvokeAgentRuntime, diff --git a/cdk/test/constructs/admission-queue-pickup.test.ts b/cdk/test/constructs/admission-queue-pickup.test.ts new file mode 100644 index 00000000..9bb3f71b --- /dev/null +++ b/cdk/test/constructs/admission-queue-pickup.test.ts @@ -0,0 +1,120 @@ +/** + * 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 { App, Stack } from 'aws-cdk-lib'; +import { Template, Match } from 'aws-cdk-lib/assertions'; +import * as dynamodb from 'aws-cdk-lib/aws-dynamodb'; +import { AdmissionQueuePickup } from '../../src/constructs/admission-queue-pickup'; + +const ORCHESTRATOR_ARN = 'arn:aws:lambda:us-east-1:123456789012:function:orchestrator:live'; + +let cachedTemplate: Template | undefined; + +function createStack(): Template { + if (cachedTemplate) return cachedTemplate; + const app = new App(); + const stack = new Stack(app, 'TestStack'); + + const taskTable = new dynamodb.Table(stack, 'TaskTable', { + partitionKey: { name: 'task_id', type: dynamodb.AttributeType.STRING }, + }); + const taskEventsTable = new dynamodb.Table(stack, 'TaskEventsTable', { + partitionKey: { name: 'task_id', type: dynamodb.AttributeType.STRING }, + sortKey: { name: 'event_id', type: dynamodb.AttributeType.STRING }, + }); + const userConcurrencyTable = new dynamodb.Table(stack, 'UserConcurrencyTable', { + partitionKey: { name: 'user_id', type: dynamodb.AttributeType.STRING }, + }); + + new AdmissionQueuePickup(stack, 'AdmissionQueuePickup', { + taskTable, + taskEventsTable, + userConcurrencyTable, + orchestratorFunctionArn: ORCHESTRATOR_ARN, + }); + + cachedTemplate = Template.fromStack(stack); + return cachedTemplate; +} + +describe('AdmissionQueuePickup construct', () => { + test('creates a Lambda function on Node 24 / ARM', () => { + const template = createStack(); + template.hasResourceProperties('AWS::Lambda::Function', { + Runtime: 'nodejs24.x', + Architectures: ['arm64'], + Timeout: 300, + }); + }); + + test('schedules the pickup every minute by default (queue latency is user-visible wait)', () => { + const template = createStack(); + template.hasResourceProperties('AWS::Events::Rule', { + ScheduleExpression: 'rate(1 minute)', + }); + }); + + test('threads queue configuration into the handler environment', () => { + const template = createStack(); + template.hasResourceProperties('AWS::Lambda::Function', { + Environment: { + Variables: Match.objectLike({ + TASK_TABLE_NAME: Match.anyValue(), + TASK_EVENTS_TABLE_NAME: Match.anyValue(), + USER_CONCURRENCY_TABLE_NAME: Match.anyValue(), + ORCHESTRATOR_FUNCTION_ARN: ORCHESTRATOR_ARN, + MAX_CONCURRENT_TASKS_PER_USER: '10', + QUEUE_MAX_AGE_SECONDS: '86400', + }), + }, + }); + }); + + test('grants lambda:InvokeFunction on the orchestrator alias only', () => { + const template = createStack(); + template.hasResourceProperties('AWS::IAM::Policy', { + PolicyDocument: { + Statement: Match.arrayWith([ + Match.objectLike({ + Action: 'lambda:InvokeFunction', + Effect: 'Allow', + Resource: ORCHESTRATOR_ARN, + }), + ]), + }, + }); + }); + + test('concurrency table access is read-only (admissionControl stays the single counter writer)', () => { + const template = createStack(); + const policies = template.findResources('AWS::IAM::Policy'); + // Collect every statement that touches the concurrency table and + // assert none of them include write actions. + const statements = Object.values(policies).flatMap( + (p: any) => p.Properties.PolicyDocument.Statement as any[], + ); + const concurrencyWrites = statements.filter(s => { + const actions: string[] = Array.isArray(s.Action) ? s.Action : [s.Action]; + const resources = JSON.stringify(s.Resource ?? ''); + return resources.includes('UserConcurrencyTable') + && actions.some(a => /PutItem|UpdateItem|DeleteItem|BatchWriteItem/.test(a)); + }); + expect(concurrencyWrites).toHaveLength(0); + }); +}); diff --git a/cdk/test/constructs/task-status.test.ts b/cdk/test/constructs/task-status.test.ts index cf01ee8c..284d3fea 100644 --- a/cdk/test/constructs/task-status.test.ts +++ b/cdk/test/constructs/task-status.test.ts @@ -22,14 +22,15 @@ import { ACTIVE_STATUSES, PRE_ACTIVE_STATUSES, TaskStatus, TaskStatusType, TERMI const ALL_STATUSES: TaskStatusType[] = Object.values(TaskStatus); describe('TaskStatus', () => { - test('defines exactly 10 states', () => { - // 8 original + AWAITING_APPROVAL (Cedar HITL gates, §10.3) + PENDING_UPLOADS (attachments). - expect(ALL_STATUSES).toHaveLength(10); + test('defines exactly 11 states', () => { + // 8 original + AWAITING_APPROVAL (Cedar HITL gates, §10.3) + // + PENDING_UPLOADS (attachments) + QUEUED (admission queue, #441). + expect(ALL_STATUSES).toHaveLength(11); }); test('contains all expected states', () => { expect(ALL_STATUSES).toEqual(expect.arrayContaining([ - 'PENDING_UPLOADS', 'SUBMITTED', 'HYDRATING', 'RUNNING', 'AWAITING_APPROVAL', 'FINALIZING', + 'PENDING_UPLOADS', 'QUEUED', 'SUBMITTED', 'HYDRATING', 'RUNNING', 'AWAITING_APPROVAL', 'FINALIZING', 'COMPLETED', 'FAILED', 'CANCELLED', 'TIMED_OUT', ])); }); @@ -43,6 +44,11 @@ describe('TaskStatus', () => { expect(TaskStatus.PENDING_UPLOADS).toBe('PENDING_UPLOADS'); expect(ALL_STATUSES).toContain('PENDING_UPLOADS'); }); + + test('QUEUED is included as a distinct state (#441)', () => { + expect(TaskStatus.QUEUED).toBe('QUEUED'); + expect(ALL_STATUSES).toContain('QUEUED'); + }); }); describe('TERMINAL_STATUSES', () => { @@ -79,21 +85,33 @@ describe('ACTIVE_STATUSES', () => { }); describe('PRE_ACTIVE_STATUSES', () => { - test('contains exactly 1 pre-active state', () => { - expect(PRE_ACTIVE_STATUSES).toHaveLength(1); + test('contains exactly 2 pre-active states', () => { + expect(PRE_ACTIVE_STATUSES).toHaveLength(2); }); - test('contains PENDING_UPLOADS', () => { + test('contains PENDING_UPLOADS and QUEUED', () => { expect(PRE_ACTIVE_STATUSES).toContain(TaskStatus.PENDING_UPLOADS); + expect(PRE_ACTIVE_STATUSES).toContain(TaskStatus.QUEUED); }); test('PENDING_UPLOADS is NOT in ACTIVE_STATUSES (no concurrency slot consumed)', () => { expect(ACTIVE_STATUSES).not.toContain(TaskStatus.PENDING_UPLOADS); }); + test('QUEUED is NOT in ACTIVE_STATUSES (#441 — counting it as active would deadlock the queue)', () => { + // The concurrency reconciler recomputes active_count from + // ACTIVE_STATUSES. If QUEUED counted as active, queued tasks would + // hold the very slots they are waiting for. + expect(ACTIVE_STATUSES).not.toContain(TaskStatus.QUEUED); + }); + test('PENDING_UPLOADS is NOT terminal', () => { expect(TERMINAL_STATUSES).not.toContain(TaskStatus.PENDING_UPLOADS); }); + + test('QUEUED is NOT terminal', () => { + expect(TERMINAL_STATUSES).not.toContain(TaskStatus.QUEUED); + }); }); describe('TERMINAL_STATUSES, ACTIVE_STATUSES, and PRE_ACTIVE_STATUSES', () => { @@ -201,4 +219,26 @@ describe('VALID_TRANSITIONS', () => { expect(VALID_TRANSITIONS[status].length).toBeGreaterThan(0); } }); + + test('SUBMITTED can transition to QUEUED (admission cap hit, #441)', () => { + expect(VALID_TRANSITIONS[TaskStatus.SUBMITTED]).toContain(TaskStatus.QUEUED); + }); + + test('QUEUED can transition back to SUBMITTED (queue pickup)', () => { + expect(VALID_TRANSITIONS[TaskStatus.QUEUED]).toContain(TaskStatus.SUBMITTED); + }); + + test('QUEUED can transition to CANCELLED (user cancel removes from queue)', () => { + expect(VALID_TRANSITIONS[TaskStatus.QUEUED]).toContain(TaskStatus.CANCELLED); + }); + + test('QUEUED can transition to FAILED (queue-age backstop)', () => { + expect(VALID_TRANSITIONS[TaskStatus.QUEUED]).toContain(TaskStatus.FAILED); + }); + + test('QUEUED cannot skip admission to RUNNING or later states', () => { + expect(VALID_TRANSITIONS[TaskStatus.QUEUED]).not.toContain(TaskStatus.HYDRATING); + expect(VALID_TRANSITIONS[TaskStatus.QUEUED]).not.toContain(TaskStatus.RUNNING); + expect(VALID_TRANSITIONS[TaskStatus.QUEUED]).not.toContain(TaskStatus.FINALIZING); + }); }); diff --git a/cdk/test/handlers/get-task.test.ts b/cdk/test/handlers/get-task.test.ts index 3a59b6e7..979f703c 100644 --- a/cdk/test/handlers/get-task.test.ts +++ b/cdk/test/handlers/get-task.test.ts @@ -25,11 +25,13 @@ jest.mock('@aws-sdk/client-dynamodb', () => ({ DynamoDBClient: jest.fn(() => ({} jest.mock('@aws-sdk/lib-dynamodb', () => ({ DynamoDBDocumentClient: { from: jest.fn(() => ({ send: mockSend })) }, GetCommand: jest.fn((input: unknown) => ({ _type: 'Get', input })), + QueryCommand: jest.fn((input: unknown) => ({ _type: 'Query', input })), })); jest.mock('ulid', () => ({ ulid: jest.fn(() => 'REQ-ULID') })); process.env.TASK_TABLE_NAME = 'Tasks'; +process.env.MAX_CONCURRENT_TASKS_PER_USER = '3'; import { handler } from '../../src/handlers/get-task'; @@ -133,6 +135,79 @@ describe('get-task handler', () => { expect(body.data.channel_source).toBe('webhook'); }); + test('non-QUEUED task carries null queue fields and issues no GSI query', async () => { + mockSend.mockReset(); + mockSend.mockResolvedValueOnce({ Item: TASK_RECORD }); + + const result = await handler(makeEvent()); + const body = JSON.parse(result.body); + expect(body.data.queue_position).toBeNull(); + expect(body.data.estimated_wait_s).toBeNull(); + expect(body.data.queued_at).toBeNull(); + expect(mockSend).toHaveBeenCalledTimes(1); // GetCommand only + }); + + test('QUEUED task returns 1-based FIFO position among the user\'s queued tasks (#441)', async () => { + mockSend.mockReset(); + mockSend + .mockResolvedValueOnce({ + Item: { + ...TASK_RECORD, + status: 'QUEUED', + status_created_at: 'QUEUED#2025-03-15T10:30:05Z', + queued_at: '2025-03-15T10:30:05Z', + }, + }) + // UserStatusIndex query: two tasks queued before ours, one after. + .mockResolvedValueOnce({ + Items: [ + { task_id: 'older-1', created_at: '2025-03-15T10:28:00Z' }, + { task_id: 'older-2', created_at: '2025-03-15T10:29:00Z' }, + { task_id: 'task-1', created_at: '2025-03-15T10:30:00Z' }, + { task_id: 'newer-1', created_at: '2025-03-15T10:31:00Z' }, + ], + }); + + const result = await handler(makeEvent()); + expect(result.statusCode).toBe(200); + const body = JSON.parse(result.body); + expect(body.data.queue_position).toBe(3); // 2 ahead + 1 + expect(body.data.estimated_wait_s).toBe(600); // ceil(3/3) * 600 + expect(body.data.queued_at).toBe('2025-03-15T10:30:05Z'); + }); + + test('QUEUED task queue-position lookup fails open to null on GSI error', async () => { + mockSend.mockReset(); + mockSend + .mockResolvedValueOnce({ + Item: { ...TASK_RECORD, status: 'QUEUED', status_created_at: 'QUEUED#2025-03-15T10:30:05Z' }, + }) + .mockRejectedValueOnce(new Error('GSI throttled')); + + const result = await handler(makeEvent()); + expect(result.statusCode).toBe(200); // GET still succeeds + const body = JSON.parse(result.body); + expect(body.data.status).toBe('QUEUED'); + expect(body.data.queue_position).toBeNull(); + expect(body.data.estimated_wait_s).toBeNull(); + }); + + test('QUEUED task that left the queue between reads reports null position (no stale rank)', async () => { + mockSend.mockReset(); + mockSend + .mockResolvedValueOnce({ + Item: { ...TASK_RECORD, status: 'QUEUED', status_created_at: 'QUEUED#2025-03-15T10:30:05Z' }, + }) + // Query result no longer contains task-1 (picked up concurrently). + .mockResolvedValueOnce({ + Items: [{ task_id: 'other', created_at: '2025-03-15T10:28:00Z' }], + }); + + const result = await handler(makeEvent()); + const body = JSON.parse(result.body); + expect(body.data.queue_position).toBeNull(); + }); + test('returns 401 when user is not authenticated', async () => { const event = makeEvent(); event.requestContext.authorizer = null; diff --git a/cdk/test/handlers/orchestrate-task.test.ts b/cdk/test/handlers/orchestrate-task.test.ts index b6ec4e44..9df671ba 100644 --- a/cdk/test/handlers/orchestrate-task.test.ts +++ b/cdk/test/handlers/orchestrate-task.test.ts @@ -76,6 +76,7 @@ import { loadBlueprintConfig, loadTask, pollTaskStatus, + queueTask, transitionTask, } from '../../src/handlers/shared/orchestrator'; @@ -908,6 +909,61 @@ describe('failTask', () => { }); }); +describe('queueTask (#441 admission queue)', () => { + test('transitions SUBMITTED -> QUEUED, stamps queued_at + admission_attempts, emits admission_queued', async () => { + mockDdbSend.mockResolvedValue({}); + const result = await queueTask(baseTask as any); + expect(result).toBe(true); + + const transition = mockDdbSend.mock.calls[0][0]; + expect(transition._type).toBe('Update'); + expect(transition.input.ExpressionAttributeValues[':fromStatus']).toBe('SUBMITTED'); + expect(transition.input.ExpressionAttributeValues[':toStatus']).toBe('QUEUED'); + // First queue entry stamps queued_at and attempts=1 + expect(transition.input.ExpressionAttributeValues[':attr_queued_at']).toBeDefined(); + expect(transition.input.ExpressionAttributeValues[':attr_admission_attempts']).toBe(1); + + const event = mockDdbSend.mock.calls[1][0]; + expect(event._type).toBe('Put'); + expect(event.input.Item.event_type).toBe('admission_queued'); + expect(event.input.Item.metadata.reason).toBe('concurrency_limit'); + }); + + test('re-queue preserves original queued_at and increments admission_attempts', async () => { + mockDdbSend.mockResolvedValue({}); + const requeued = { ...baseTask, queued_at: '2024-01-01T00:00:00Z', admission_attempts: 2 }; + await queueTask(requeued as any); + + const transition = mockDdbSend.mock.calls[0][0]; + // queued_at already set — must NOT be overwritten + expect(transition.input.ExpressionAttributeValues[':attr_queued_at']).toBeUndefined(); + expect(transition.input.ExpressionAttributeValues[':attr_admission_attempts']).toBe(3); + }); + + test('returns false without emitting when a concurrent transition wins the conditional check', async () => { + const condErr = new Error('The conditional request failed'); + condErr.name = 'ConditionalCheckFailedException'; + mockDdbSend.mockRejectedValueOnce(condErr); + + const result = await queueTask(baseTask as any); + expect(result).toBe(false); + expect(mockDdbSend).toHaveBeenCalledTimes(1); // transition attempt only, no event + }); + + test('rethrows unexpected DDB errors so the durable step retries', async () => { + mockDdbSend.mockRejectedValueOnce(new Error('DDB timeout')); + await expect(queueTask(baseTask as any)).rejects.toThrow('DDB timeout'); + }); + + test('never touches the concurrency counter (no slot held while QUEUED)', async () => { + mockDdbSend.mockResolvedValue({}); + await queueTask(baseTask as any); + for (const call of mockDdbSend.mock.calls) { + expect(call[0].input.TableName).not.toBe('UserConcurrency'); + } + }); +}); + describe('emitTaskEvent', () => { test('writes event to DynamoDB with TTL', async () => { mockDdbSend.mockResolvedValueOnce({}); diff --git a/cdk/test/handlers/reconcile-admission-queue.test.ts b/cdk/test/handlers/reconcile-admission-queue.test.ts new file mode 100644 index 00000000..9a824358 --- /dev/null +++ b/cdk/test/handlers/reconcile-admission-queue.test.ts @@ -0,0 +1,404 @@ +/** + * 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. + */ + +// Admission-queue deferred pickup (#441). Regression coverage for the #331 +// scenario lives in the "fan-out burst" describe block at the bottom. + +// --- Mocks --- +const mockDdbSend = jest.fn(); +jest.mock('@aws-sdk/client-dynamodb', () => ({ + DynamoDBClient: jest.fn(() => ({ send: mockDdbSend })), + QueryCommand: jest.fn((input: unknown) => ({ _type: 'Query', input })), + GetItemCommand: jest.fn((input: unknown) => ({ _type: 'GetItem', input })), + UpdateItemCommand: jest.fn((input: unknown) => ({ _type: 'UpdateItem', input })), + PutItemCommand: jest.fn((input: unknown) => ({ _type: 'PutItem', input })), +})); + +const mockLambdaSend = jest.fn(); +jest.mock('@aws-sdk/client-lambda', () => ({ + LambdaClient: jest.fn(() => ({ send: mockLambdaSend })), + InvokeCommand: jest.fn((input: unknown) => ({ _type: 'Invoke', input })), +})); + +let ulidCounter = 0; +jest.mock('ulid', () => ({ ulid: jest.fn(() => `ULID${ulidCounter++}`) })); + +process.env.TASK_TABLE_NAME = 'Tasks'; +process.env.TASK_EVENTS_TABLE_NAME = 'TaskEvents'; +process.env.USER_CONCURRENCY_TABLE_NAME = 'Concurrency'; +process.env.ORCHESTRATOR_FUNCTION_ARN = 'arn:aws:lambda:us-east-1:123456789012:function:orchestrator:live'; +process.env.MAX_CONCURRENT_TASKS_PER_USER = '3'; +process.env.QUEUE_MAX_AGE_SECONDS = '86400'; +process.env.TASK_RETENTION_DAYS = '90'; + +import { handler } from '../../src/handlers/reconcile-admission-queue'; + +/** A QUEUED StatusIndex row as the DDB low-level client returns it. */ +function queuedRow(opts: { + task_id: string; + user_id: string; + created_at: string; + admission_attempts?: number; +}): Record { + return { + task_id: { S: opts.task_id }, + user_id: { S: opts.user_id }, + created_at: { S: opts.created_at }, + ...(opts.admission_attempts !== undefined && { admission_attempts: { N: String(opts.admission_attempts) } }), + }; +} + +/** Recent ISO timestamp (now - ageSeconds). */ +function isoAge(ageSeconds: number): string { + return new Date(Date.now() - ageSeconds * 1000).toISOString(); +} + +interface SentCommand { + _type: string; + input: Record; +} + +function sentDdbCommands(type?: string): SentCommand[] { + const cmds = mockDdbSend.mock.calls.map(c => c[0] as SentCommand); + return type ? cmds.filter(c => c._type === type) : cmds; +} + +beforeEach(() => { + jest.clearAllMocks(); + ulidCounter = 0; +}); + +describe('reconcile-admission-queue — empty queue', () => { + test('no QUEUED tasks → single query, no writes, no invokes', async () => { + mockDdbSend.mockResolvedValueOnce({ Items: [] }); + const summary = await handler(); + expect(summary.queued_seen).toBe(0); + expect(summary.picked_up).toBe(0); + expect(mockDdbSend).toHaveBeenCalledTimes(1); + expect(mockLambdaSend).not.toHaveBeenCalled(); + }); +}); + +describe('reconcile-admission-queue — pickup', () => { + test('picks up the oldest task when the user has capacity', async () => { + mockDdbSend.mockImplementation((cmd: SentCommand) => { + if (cmd._type === 'Query') { + return Promise.resolve({ + Items: [queuedRow({ task_id: 'T1', user_id: 'u1', created_at: isoAge(60), admission_attempts: 1 })], + }); + } + if (cmd._type === 'GetItem') { + return Promise.resolve({ Item: { active_count: { N: '1' } } }); // 1 of 3 used + } + return Promise.resolve({}); + }); + mockLambdaSend.mockResolvedValue({}); + + const summary = await handler(); + expect(summary.picked_up).toBe(1); + + // Conditional QUEUED -> SUBMITTED flip + const updates = sentDdbCommands('UpdateItem'); + expect(updates).toHaveLength(1); + expect(updates[0].input.ConditionExpression).toContain(':queued'); + expect(updates[0].input.ExpressionAttributeValues[':submitted']).toEqual({ S: 'SUBMITTED' }); + + // queue_pickup event emitted + const events = sentDdbCommands('PutItem'); + expect(events).toHaveLength(1); + expect(events[0].input.Item.event_type).toEqual({ S: 'queue_pickup' }); + + // Orchestrator re-invoked async with a pickup nonce + expect(mockLambdaSend).toHaveBeenCalledTimes(1); + const invoke = mockLambdaSend.mock.calls[0][0] as SentCommand; + expect(invoke.input.InvocationType).toBe('Event'); + const payload = JSON.parse(new TextDecoder().decode(invoke.input.Payload)); + expect(payload.task_id).toBe('T1'); + expect(payload.queue_pickup_id).toBeDefined(); + }); + + test('picks up in FIFO order, bounded by free capacity', async () => { + // 3 queued tasks for u1; only 2 free slots (active_count=1, max=3). + mockDdbSend.mockImplementation((cmd: SentCommand) => { + if (cmd._type === 'Query') { + return Promise.resolve({ + Items: [ + queuedRow({ task_id: 'OLD', user_id: 'u1', created_at: isoAge(300) }), + queuedRow({ task_id: 'MID', user_id: 'u1', created_at: isoAge(240) }), + queuedRow({ task_id: 'NEW', user_id: 'u1', created_at: isoAge(180) }), + ], + }); + } + if (cmd._type === 'GetItem') { + return Promise.resolve({ Item: { active_count: { N: '1' } } }); + } + return Promise.resolve({}); + }); + mockLambdaSend.mockResolvedValue({}); + + const summary = await handler(); + expect(summary.picked_up).toBe(2); + expect(summary.skipped_no_capacity).toBe(1); + + const flipped = sentDdbCommands('UpdateItem').map(u => u.input.Key.task_id.S); + expect(flipped).toEqual(['OLD', 'MID']); // FIFO — NEW stays queued + }); + + test('skips a user at capacity entirely', async () => { + mockDdbSend.mockImplementation((cmd: SentCommand) => { + if (cmd._type === 'Query') { + return Promise.resolve({ + Items: [queuedRow({ task_id: 'T1', user_id: 'u1', created_at: isoAge(60) })], + }); + } + if (cmd._type === 'GetItem') { + return Promise.resolve({ Item: { active_count: { N: '3' } } }); // full + } + return Promise.resolve({}); + }); + + const summary = await handler(); + expect(summary.picked_up).toBe(0); + expect(summary.skipped_no_capacity).toBe(1); + expect(sentDdbCommands('UpdateItem')).toHaveLength(0); + expect(mockLambdaSend).not.toHaveBeenCalled(); + }); + + test('per-user isolation: one user at capacity does not block another', async () => { + mockDdbSend.mockImplementation((cmd: SentCommand) => { + if (cmd._type === 'Query') { + return Promise.resolve({ + Items: [ + queuedRow({ task_id: 'FULL-USER-TASK', user_id: 'u-full', created_at: isoAge(120) }), + queuedRow({ task_id: 'FREE-USER-TASK', user_id: 'u-free', created_at: isoAge(60) }), + ], + }); + } + if (cmd._type === 'GetItem') { + const userId = cmd.input.Key.user_id.S; + return Promise.resolve({ Item: { active_count: { N: userId === 'u-full' ? '3' : '0' } } }); + } + return Promise.resolve({}); + }); + mockLambdaSend.mockResolvedValue({}); + + const summary = await handler(); + expect(summary.picked_up).toBe(1); + expect(summary.skipped_no_capacity).toBe(1); + const flipped = sentDdbCommands('UpdateItem').map(u => u.input.Key.task_id.S); + expect(flipped).toEqual(['FREE-USER-TASK']); + }); + + test('missing concurrency row counts as zero active (full capacity)', async () => { + mockDdbSend.mockImplementation((cmd: SentCommand) => { + if (cmd._type === 'Query') { + return Promise.resolve({ + Items: [queuedRow({ task_id: 'T1', user_id: 'new-user', created_at: isoAge(30) })], + }); + } + if (cmd._type === 'GetItem') { + return Promise.resolve({}); // no Item + } + return Promise.resolve({}); + }); + mockLambdaSend.mockResolvedValue({}); + + const summary = await handler(); + expect(summary.picked_up).toBe(1); + }); +}); + +describe('reconcile-admission-queue — races and failures', () => { + test('lost flip race (ConditionalCheckFailedException) is skipped cleanly', async () => { + const condErr = Object.assign(new Error('conditional failed'), { name: 'ConditionalCheckFailedException' }); + mockDdbSend.mockImplementation((cmd: SentCommand) => { + if (cmd._type === 'Query') { + return Promise.resolve({ + Items: [queuedRow({ task_id: 'T1', user_id: 'u1', created_at: isoAge(60) })], + }); + } + if (cmd._type === 'GetItem') { + return Promise.resolve({ Item: { active_count: { N: '0' } } }); + } + if (cmd._type === 'UpdateItem') { + return Promise.reject(condErr); + } + return Promise.resolve({}); + }); + + const summary = await handler(); + expect(summary.picked_up).toBe(0); + expect(summary.skipped_race).toBe(1); + expect(summary.errors).toBe(0); + expect(mockLambdaSend).not.toHaveBeenCalled(); + }); + + test('orchestrator invoke failure after flip does NOT roll back (stranded reconciler sweeps)', async () => { + mockDdbSend.mockImplementation((cmd: SentCommand) => { + if (cmd._type === 'Query') { + return Promise.resolve({ + Items: [queuedRow({ task_id: 'T1', user_id: 'u1', created_at: isoAge(60) })], + }); + } + if (cmd._type === 'GetItem') { + return Promise.resolve({ Item: { active_count: { N: '0' } } }); + } + return Promise.resolve({}); + }); + mockLambdaSend.mockRejectedValue(new Error('lambda throttled')); + + const summary = await handler(); + expect(summary.picked_up).toBe(0); + // Only the QUEUED -> SUBMITTED flip — no compensating write back to QUEUED. + const updates = sentDdbCommands('UpdateItem'); + expect(updates).toHaveLength(1); + expect(updates[0].input.ExpressionAttributeValues[':submitted']).toEqual({ S: 'SUBMITTED' }); + }); + + test('per-user concurrency read failure skips that user but not others', async () => { + mockDdbSend.mockImplementation((cmd: SentCommand) => { + if (cmd._type === 'Query') { + return Promise.resolve({ + Items: [ + queuedRow({ task_id: 'BROKEN', user_id: 'u-err', created_at: isoAge(120) }), + queuedRow({ task_id: 'OK', user_id: 'u-ok', created_at: isoAge(60) }), + ], + }); + } + if (cmd._type === 'GetItem') { + if (cmd.input.Key.user_id.S === 'u-err') { + return Promise.reject(new Error('DDB throttled')); + } + return Promise.resolve({ Item: { active_count: { N: '0' } } }); + } + return Promise.resolve({}); + }); + mockLambdaSend.mockResolvedValue({}); + + const summary = await handler(); + expect(summary.errors).toBe(1); + expect(summary.picked_up).toBe(1); + const flipped = sentDdbCommands('UpdateItem').map(u => u.input.Key.task_id.S); + expect(flipped).toEqual(['OK']); + }); + + test('queue query failure aborts the cycle with a throw (EventBridge will retry)', async () => { + mockDdbSend.mockRejectedValueOnce(new Error('GSI unavailable')); + await expect(handler()).rejects.toThrow('GSI unavailable'); + }); +}); + +describe('reconcile-admission-queue — max-age backstop', () => { + test('a task queued past QUEUE_MAX_AGE_SECONDS is failed, not picked up', async () => { + mockDdbSend.mockImplementation((cmd: SentCommand) => { + if (cmd._type === 'Query') { + return Promise.resolve({ + Items: [queuedRow({ task_id: 'ZOMBIE', user_id: 'u1', created_at: isoAge(90000) })], // > 86400 + }); + } + if (cmd._type === 'GetItem') { + return Promise.resolve({ Item: { active_count: { N: '0' } } }); + } + return Promise.resolve({}); + }); + + const summary = await handler(); + expect(summary.expired).toBe(1); + expect(summary.picked_up).toBe(0); + + const updates = sentDdbCommands('UpdateItem'); + expect(updates).toHaveLength(1); + expect(updates[0].input.ExpressionAttributeValues[':failed']).toEqual({ S: 'FAILED' }); + // No concurrency decrement — QUEUED tasks never held a slot. + expect(updates.every(u => u.input.TableName === 'Tasks')).toBe(true); + // task_failed event with queue_timeout reason + const events = sentDdbCommands('PutItem'); + expect(events).toHaveLength(1); + expect(events[0].input.Item.event_type).toEqual({ S: 'task_failed' }); + expect(events[0].input.Item.metadata.M.reason).toEqual({ S: 'queue_timeout' }); + expect(mockLambdaSend).not.toHaveBeenCalled(); + }); +}); + +describe('reconcile-admission-queue — #331 regression (fan-out burst)', () => { + test('a burst of children above the cap all survive as QUEUED and drain FIFO without any failure', async () => { + // #331: an epic fan-out released 8 children against a cap of 3. + // Pre-#441 admission FAILED the excess 5. Now: 3 admitted directly + // (never reach QUEUED), 5 QUEUED. This cycle: user has 0 active + // (children finished), so 3 of the 5 queued drain; 2 remain queued; + // ZERO tasks fail. + const queued = ['C4', 'C5', 'C6', 'C7', 'C8'].map((id, i) => + // Oldest first: C4 queued 300s ago ... C8 queued 60s ago. + queuedRow({ task_id: id, user_id: 'epic-user', created_at: isoAge(300 - i * 60), admission_attempts: 1 }), + ); + mockDdbSend.mockImplementation((cmd: SentCommand) => { + if (cmd._type === 'Query') { + return Promise.resolve({ Items: queued }); + } + if (cmd._type === 'GetItem') { + return Promise.resolve({ Item: { active_count: { N: '0' } } }); + } + return Promise.resolve({}); + }); + mockLambdaSend.mockResolvedValue({}); + + const summary = await handler(); + expect(summary.queued_seen).toBe(5); + expect(summary.picked_up).toBe(3); // cap-bounded drain + expect(summary.skipped_no_capacity).toBe(2); // survive for next cycle + expect(summary.expired).toBe(0); // nothing dies + + // FIFO: the three oldest children drain first. + const flipped = sentDdbCommands('UpdateItem').map(u => u.input.Key.task_id.S); + expect(flipped).toEqual(['C4', 'C5', 'C6']); + // No task was transitioned to FAILED anywhere in the cycle. + const failedWrites = sentDdbCommands('UpdateItem') + .filter(u => u.input.ExpressionAttributeValues[':failed'] !== undefined); + expect(failedWrites).toHaveLength(0); + }); + + test('paginated QUEUED query preserves FIFO across pages', async () => { + let queryCount = 0; + mockDdbSend.mockImplementation((cmd: SentCommand) => { + if (cmd._type === 'Query') { + queryCount++; + if (queryCount === 1) { + return Promise.resolve({ + Items: [queuedRow({ task_id: 'P1', user_id: 'u1', created_at: isoAge(120) })], + LastEvaluatedKey: { task_id: { S: 'P1' } }, + }); + } + return Promise.resolve({ + Items: [queuedRow({ task_id: 'P2', user_id: 'u1', created_at: isoAge(60) })], + }); + } + if (cmd._type === 'GetItem') { + return Promise.resolve({ Item: { active_count: { N: '1' } } }); // 2 free + } + return Promise.resolve({}); + }); + mockLambdaSend.mockResolvedValue({}); + + const summary = await handler(); + expect(summary.queued_seen).toBe(2); + expect(summary.picked_up).toBe(2); + const flipped = sentDdbCommands('UpdateItem').map(u => u.input.Key.task_id.S); + expect(flipped).toEqual(['P1', 'P2']); + }); +}); diff --git a/cli/src/format.ts b/cli/src/format.ts index 1ebb2240..cd69e299 100644 --- a/cli/src/format.ts +++ b/cli/src/format.ts @@ -56,6 +56,10 @@ export function formatTaskDetail(task: TaskDetail): string { if (task.branch_name) { lines.push(`Branch: ${task.branch_name}`); } + // Admission queue (#441): show FIFO position + rough ETA while QUEUED. + if (task.status === 'QUEUED' && task.queue_position !== null) { + lines.push(`Queue: position ${task.queue_position}${formatQueueEta(task.estimated_wait_s)}`); + } if (task.max_turns !== null) { lines.push(`Max Turns: ${task.max_turns}`); } @@ -96,6 +100,16 @@ export function formatTaskDetail(task: TaskDetail): string { return lines.join('\n'); } +/** Render the ``~Xm wait`` suffix for a queued task's position line (#441). + * Empty string when the server could not estimate a wait. */ +function formatQueueEta(estimatedWaitS: number | null): string { + if (estimatedWaitS === null || !Number.isFinite(estimatedWaitS)) { + return ''; + } + const minutes = Math.max(1, Math.round(estimatedWaitS / 60)); + return ` (est. wait ~${minutes}m)`; +} + /** Format a list of TaskSummary as an aligned table. */ export function formatTaskList(tasks: TaskSummary[]): string { if (tasks.length === 0) { @@ -195,6 +209,11 @@ export function formatStatusSnapshot( if (task.task_description) { lines.push(...formatDescriptionLines(task.task_description)); } + // Admission queue (#441): while QUEUED the pipeline has not started, so + // position + ETA are the most useful lines a user can see. + if (task.status === 'QUEUED' && task.queue_position !== null) { + lines.push(` Queue: position ${task.queue_position}${formatQueueEta(task.estimated_wait_s)}`); + } lines.push( ` Turn: ${describeTurn(task, lastTurnEvent)}`, ` Last milestone: ${describeMilestone(milestoneEvent, now)}`, diff --git a/cli/src/types.ts b/cli/src/types.ts index 7ae2c633..76a91380 100644 --- a/cli/src/types.ts +++ b/cli/src/types.ts @@ -41,6 +41,7 @@ export type AttachmentType = 'image' | 'file' | 'url'; */ export type TaskStatusType = | 'PENDING_UPLOADS' + | 'QUEUED' | 'SUBMITTED' | 'HYDRATING' | 'RUNNING' @@ -151,6 +152,18 @@ export interface TaskDetail { /** Cedar HITL: when ``status = AWAITING_APPROVAL``, the * ``request_id`` of the pending approval row. Null otherwise. */ readonly awaiting_approval_request_id: string | null; + /** Admission queue (#441): ISO timestamp the task first entered + * QUEUED; null for tasks that were admitted directly. Mirrors + * ``cdk/src/handlers/shared/types.ts::TaskDetail``. */ + readonly queued_at: string | null; + /** Admission queue (#441): 1-based FIFO position among the caller's + * QUEUED tasks when ``status = QUEUED``; null otherwise. Computed + * at read time by the server — it changes as the queue drains. */ + readonly queue_position: number | null; + /** Admission queue (#441): rough ETA (seconds) until pickup, derived + * from queue position and recent task durations. Null when + * ``queue_position`` is null. */ + readonly estimated_wait_s: number | null; } /** Response body of ``GET /v1/tasks/{task_id}/trace`` (design §10.1). */ diff --git a/cli/test/format-status-snapshot.test.ts b/cli/test/format-status-snapshot.test.ts index ea72cb89..e1d840d7 100644 --- a/cli/test/format-status-snapshot.test.ts +++ b/cli/test/format-status-snapshot.test.ts @@ -60,6 +60,9 @@ function buildTask(overrides: Partial = {}): TaskDetail { approval_gate_count: 0, approval_gate_cap: 50, awaiting_approval_request_id: null, + queued_at: null, + queue_position: null, + estimated_wait_s: null, ...overrides, }; } @@ -75,6 +78,25 @@ function mkEvent(overrides: Partial): TaskEvent { } describe('formatStatusSnapshot', () => { + test('QUEUED task shows queue position + ETA line (#441)', () => { + const task = buildTask({ + status: 'QUEUED', + started_at: null, + queued_at: '2026-04-29T15:27:30Z', + queue_position: 2, + estimated_wait_s: 600, + }); + const rendered = formatStatusSnapshot(task, [], NOW); + expect(rendered).toContain('Task abc123 — QUEUED'); + expect(rendered).toContain(' Queue: position 2 (est. wait ~10m)'); + }); + + test('QUEUED task with null position omits the queue line (GSI fail-open)', () => { + const task = buildTask({ status: 'QUEUED', started_at: null, queue_position: null }); + const rendered = formatStatusSnapshot(task, [], NOW); + expect(rendered).not.toContain('Queue:'); + }); + test('happy path renders the full template', () => { const task = buildTask(); // Events are newest-first per the ``?desc=1`` contract. ULIDs are diff --git a/cli/test/format.test.ts b/cli/test/format.test.ts index cf82bf8c..5d66e3d9 100644 --- a/cli/test/format.test.ts +++ b/cli/test/format.test.ts @@ -54,6 +54,9 @@ describe('format', () => { approval_gate_count: 0, approval_gate_cap: 50, awaiting_approval_request_id: null, + queued_at: null, + queue_position: null, + estimated_wait_s: null, }; describe('formatTaskDetail', () => { @@ -72,6 +75,39 @@ describe('format', () => { expect(output).toContain('Max Turns: 100'); }); + test('renders queue position + ETA for a QUEUED task (#441)', () => { + const queued: TaskDetail = { + ...task, + status: 'QUEUED', + queued_at: '2026-01-01T00:00:30Z', + queue_position: 3, + estimated_wait_s: 600, + }; + const output = formatTaskDetail(queued); + expect(output).toContain('Status: QUEUED'); + expect(output).toContain('Queue: position 3 (est. wait ~10m)'); + }); + + test('renders queue position without ETA when estimated_wait_s is null', () => { + const queued: TaskDetail = { + ...task, + status: 'QUEUED', + queue_position: 1, + estimated_wait_s: null, + }; + const output = formatTaskDetail(queued); + expect(output).toContain('Queue: position 1'); + expect(output).not.toContain('est. wait'); + }); + + test('omits the queue line when position is null (GSI fail-open) or task not QUEUED', () => { + const queuedNoPos: TaskDetail = { ...task, status: 'QUEUED', queue_position: null }; + expect(formatTaskDetail(queuedNoPos)).not.toContain('Queue:'); + // RUNNING task with a stale position value should not render it either. + const running: TaskDetail = { ...task, status: 'RUNNING', queue_position: 2 }; + expect(formatTaskDetail(running)).not.toContain('Queue:'); + }); + test('omits null fields', () => { const minimal: TaskDetail = { ...task, diff --git a/docs/design/ORCHESTRATOR.md b/docs/design/ORCHESTRATOR.md index 5103ed7e..a3c87584 100644 --- a/docs/design/ORCHESTRATOR.md +++ b/docs/design/ORCHESTRATOR.md @@ -46,13 +46,14 @@ The orchestrator is deliberately scoped. It handles coordination and bookkeeping ## Task state machine -Every task moves through a finite set of states from creation to a terminal outcome. The state machine is the backbone of the orchestrator: it determines what actions are valid at each point, when resources are acquired or released, and how the platform recovers from failures. Four of the ten states are terminal, meaning the task is done and no further transitions occur. +Every task moves through a finite set of states from creation to a terminal outcome. The state machine is the backbone of the orchestrator: it determines what actions are valid at each point, when resources are acquired or released, and how the platform recovers from failures. Four of the eleven states are terminal, meaning the task is done and no further transitions occur. ### States | State | Description | Duration | |---|---|---| | `PENDING_UPLOADS` | Presigned-upload task awaiting client file uploads | Minutes (30-min auto-cancel) | +| `QUEUED` | Admission-capped task awaiting a free concurrency slot (#441) | Minutes to hours (24h backstop) | | `SUBMITTED` | Task accepted, awaiting orchestration | Milliseconds | | `HYDRATING` | Fetching GitHub data, querying memory, assembling prompt | Seconds | | `RUNNING` | Agent session active in compute environment | Minutes to hours | @@ -74,9 +75,14 @@ stateDiagram-v2 PENDING_UPLOADS --> CANCELLED : User cancels or 30-min auto-cancel SUBMITTED --> HYDRATING : Admission passes + SUBMITTED --> QUEUED : Concurrency cap hit SUBMITTED --> FAILED : Admission rejected SUBMITTED --> CANCELLED : User cancels + QUEUED --> SUBMITTED : Queue pickup (slot free) + QUEUED --> CANCELLED : User cancels + QUEUED --> FAILED : Queue-age backstop (24h) + HYDRATING --> RUNNING : Session started HYDRATING --> AWAITING_APPROVAL : Cedar soft-deny gate HYDRATING --> FAILED : Hydration error @@ -105,7 +111,11 @@ stateDiagram-v2 | `PENDING_UPLOADS` | `FAILED` | Screening blocked | Any attachment fails security screening | | `PENDING_UPLOADS` | `CANCELLED` | User cancels or auto-cancel | Upload window expired (30 min) or explicit cancel | | `SUBMITTED` | `HYDRATING` | Admission passes | Concurrency slot acquired | -| `SUBMITTED` | `FAILED` | Admission rejected | Repo not onboarded, rate/concurrency limit, validation error | +| `SUBMITTED` | `QUEUED` | Concurrency cap hit | Per-user cap reached; task parked in FIFO admission queue (#441). No slot held. | +| `SUBMITTED` | `FAILED` | Admission rejected | Repo not onboarded, validation error | +| `QUEUED` | `SUBMITTED` | Queue pickup | Admission-queue pickup Lambda sees free capacity; re-invokes the orchestrator. FIFO by `created_at`. | +| `QUEUED` | `CANCELLED` | User cancels | Explicit cancel removes the task from the queue | +| `QUEUED` | `FAILED` | Queue-age backstop | Task waited longer than `QUEUE_MAX_AGE_SECONDS` (default 24h) without admission | | `HYDRATING` | `RUNNING` | Hydration complete | `invoke_agent_runtime` returns session ID | | `HYDRATING` | `AWAITING_APPROVAL` | Cedar soft-deny gate fires | Tool call triggers a soft-deny policy rule during hydration | | `HYDRATING` | `FAILED` | Hydration error | GitHub API failure, guardrail blocks content, Bedrock unavailable | @@ -126,6 +136,7 @@ Users can cancel a task at any point. The orchestrator's response depends on how | State when cancel arrives | Action | |---|---| | `PENDING_UPLOADS` | Transition to `CANCELLED`. Clean up S3 objects under the task's attachment prefix. No concurrency slot to release. | +| `QUEUED` | Transition to `CANCELLED`. Removes the task from the admission queue. No compute or concurrency slot to release (a queued task never held one). | | `SUBMITTED` | Transition to `CANCELLED`. No cleanup needed. | | `HYDRATING` | Abort hydration, release concurrency slot, transition to `CANCELLED`. | | `RUNNING` | Call `stop_runtime_session`, wait for confirmation, release concurrency, transition to `CANCELLED`. Partial work on GitHub remains for the user to inspect. | @@ -165,9 +176,9 @@ The orchestrator (`orchestrate-task.ts`) runs these as distinct durable-executio Validates the task before any compute is consumed. Checks run in order: 1. **Repo onboarding** - `GetItem` on `RepoTable`. If not found or inactive, reject with `REPO_NOT_ONBOARDED`. This runs at the API handler level (`createTaskCore`) for fast rejection. -2. **User concurrency** - Atomic check-and-increment on `UserConcurrency` counter. If at limit (default 3-5), reject. +2. **User concurrency** - Atomic check-and-increment on `UserConcurrency` counter. If at limit (default 10), the task is **queued, not failed** (#441): it transitions `SUBMITTED → QUEUED` and a scheduled admission-queue pickup Lambda re-attempts admission in FIFO order (by `created_at`) as slots free up, flipping `QUEUED → SUBMITTED` and re-invoking the orchestrator. The pickup Lambda does a read-only capacity pre-check; the orchestrator's atomic increment remains the single writer of the counter, so a pickup that loses the race harmlessly re-queues without losing FIFO position. `GET /tasks/{id}` surfaces `queue_position` and `estimated_wait_s` while queued. 3. **System concurrency** - Compare total running + hydrating tasks to system limit (bounded by AgentCore quotas). -4. **Rate limiting** - Sliding window counter (10 tasks/hour per user). Exceeded tasks are rejected, not queued. +4. **Rate limiting** - Sliding window counter (10 tasks/hour per user). Rate-limit rejections happen at submit time and are rejected, not queued (unlike the concurrency cap, which queues). 5. **Idempotency** - If the request includes an idempotency key and a task with that key exists, return the existing task. On acceptance, the concurrency slot is acquired and the orchestrator proceeds to pre-flight. diff --git a/docs/src/content/docs/architecture/Orchestrator.md b/docs/src/content/docs/architecture/Orchestrator.md index fd48f796..0067ba18 100644 --- a/docs/src/content/docs/architecture/Orchestrator.md +++ b/docs/src/content/docs/architecture/Orchestrator.md @@ -50,13 +50,14 @@ The orchestrator is deliberately scoped. It handles coordination and bookkeeping ## Task state machine -Every task moves through a finite set of states from creation to a terminal outcome. The state machine is the backbone of the orchestrator: it determines what actions are valid at each point, when resources are acquired or released, and how the platform recovers from failures. Four of the ten states are terminal, meaning the task is done and no further transitions occur. +Every task moves through a finite set of states from creation to a terminal outcome. The state machine is the backbone of the orchestrator: it determines what actions are valid at each point, when resources are acquired or released, and how the platform recovers from failures. Four of the eleven states are terminal, meaning the task is done and no further transitions occur. ### States | State | Description | Duration | |---|---|---| | `PENDING_UPLOADS` | Presigned-upload task awaiting client file uploads | Minutes (30-min auto-cancel) | +| `QUEUED` | Admission-capped task awaiting a free concurrency slot (#441) | Minutes to hours (24h backstop) | | `SUBMITTED` | Task accepted, awaiting orchestration | Milliseconds | | `HYDRATING` | Fetching GitHub data, querying memory, assembling prompt | Seconds | | `RUNNING` | Agent session active in compute environment | Minutes to hours | @@ -78,9 +79,14 @@ stateDiagram-v2 PENDING_UPLOADS --> CANCELLED : User cancels or 30-min auto-cancel SUBMITTED --> HYDRATING : Admission passes + SUBMITTED --> QUEUED : Concurrency cap hit SUBMITTED --> FAILED : Admission rejected SUBMITTED --> CANCELLED : User cancels + QUEUED --> SUBMITTED : Queue pickup (slot free) + QUEUED --> CANCELLED : User cancels + QUEUED --> FAILED : Queue-age backstop (24h) + HYDRATING --> RUNNING : Session started HYDRATING --> AWAITING_APPROVAL : Cedar soft-deny gate HYDRATING --> FAILED : Hydration error @@ -109,7 +115,11 @@ stateDiagram-v2 | `PENDING_UPLOADS` | `FAILED` | Screening blocked | Any attachment fails security screening | | `PENDING_UPLOADS` | `CANCELLED` | User cancels or auto-cancel | Upload window expired (30 min) or explicit cancel | | `SUBMITTED` | `HYDRATING` | Admission passes | Concurrency slot acquired | -| `SUBMITTED` | `FAILED` | Admission rejected | Repo not onboarded, rate/concurrency limit, validation error | +| `SUBMITTED` | `QUEUED` | Concurrency cap hit | Per-user cap reached; task parked in FIFO admission queue (#441). No slot held. | +| `SUBMITTED` | `FAILED` | Admission rejected | Repo not onboarded, validation error | +| `QUEUED` | `SUBMITTED` | Queue pickup | Admission-queue pickup Lambda sees free capacity; re-invokes the orchestrator. FIFO by `created_at`. | +| `QUEUED` | `CANCELLED` | User cancels | Explicit cancel removes the task from the queue | +| `QUEUED` | `FAILED` | Queue-age backstop | Task waited longer than `QUEUE_MAX_AGE_SECONDS` (default 24h) without admission | | `HYDRATING` | `RUNNING` | Hydration complete | `invoke_agent_runtime` returns session ID | | `HYDRATING` | `AWAITING_APPROVAL` | Cedar soft-deny gate fires | Tool call triggers a soft-deny policy rule during hydration | | `HYDRATING` | `FAILED` | Hydration error | GitHub API failure, guardrail blocks content, Bedrock unavailable | @@ -130,6 +140,7 @@ Users can cancel a task at any point. The orchestrator's response depends on how | State when cancel arrives | Action | |---|---| | `PENDING_UPLOADS` | Transition to `CANCELLED`. Clean up S3 objects under the task's attachment prefix. No concurrency slot to release. | +| `QUEUED` | Transition to `CANCELLED`. Removes the task from the admission queue. No compute or concurrency slot to release (a queued task never held one). | | `SUBMITTED` | Transition to `CANCELLED`. No cleanup needed. | | `HYDRATING` | Abort hydration, release concurrency slot, transition to `CANCELLED`. | | `RUNNING` | Call `stop_runtime_session`, wait for confirmation, release concurrency, transition to `CANCELLED`. Partial work on GitHub remains for the user to inspect. | @@ -169,9 +180,9 @@ The orchestrator (`orchestrate-task.ts`) runs these as distinct durable-executio Validates the task before any compute is consumed. Checks run in order: 1. **Repo onboarding** - `GetItem` on `RepoTable`. If not found or inactive, reject with `REPO_NOT_ONBOARDED`. This runs at the API handler level (`createTaskCore`) for fast rejection. -2. **User concurrency** - Atomic check-and-increment on `UserConcurrency` counter. If at limit (default 3-5), reject. +2. **User concurrency** - Atomic check-and-increment on `UserConcurrency` counter. If at limit (default 10), the task is **queued, not failed** (#441): it transitions `SUBMITTED → QUEUED` and a scheduled admission-queue pickup Lambda re-attempts admission in FIFO order (by `created_at`) as slots free up, flipping `QUEUED → SUBMITTED` and re-invoking the orchestrator. The pickup Lambda does a read-only capacity pre-check; the orchestrator's atomic increment remains the single writer of the counter, so a pickup that loses the race harmlessly re-queues without losing FIFO position. `GET /tasks/{id}` surfaces `queue_position` and `estimated_wait_s` while queued. 3. **System concurrency** - Compare total running + hydrating tasks to system limit (bounded by AgentCore quotas). -4. **Rate limiting** - Sliding window counter (10 tasks/hour per user). Exceeded tasks are rejected, not queued. +4. **Rate limiting** - Sliding window counter (10 tasks/hour per user). Rate-limit rejections happen at submit time and are rejected, not queued (unlike the concurrency cap, which queues). 5. **Idempotency** - If the request includes an idempotency key and a task with that key exists, return the existing task. On acceptance, the concurrency slot is acquired and the orchestrator proceeds to pre-flight. From fb9c232054396b11fa7a785bef95c5479a0509a4 Mon Sep 17 00:00:00 2001 From: Sphia Sadek Date: Wed, 22 Jul 2026 19:00:31 +0100 Subject: [PATCH 2/2] Update cdk/src/handlers/get-task.ts Co-authored-by: Scott Schreckengaust --- cdk/src/handlers/get-task.ts | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/cdk/src/handlers/get-task.ts b/cdk/src/handlers/get-task.ts index 488081c7..b1ea4e12 100644 --- a/cdk/src/handlers/get-task.ts +++ b/cdk/src/handlers/get-task.ts @@ -76,7 +76,13 @@ async function computeQueueInfo( 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) { + } else if ( + typeof item.created_at === 'string' + && (item.created_at < record.created_at + || (item.created_at === record.created_at + && typeof item.task_id === 'string' + && item.task_id < record.task_id)) + ) { ahead++; } }