diff --git a/.changeset/queue-combined-concurrency-stats.md b/.changeset/queue-combined-concurrency-stats.md new file mode 100644 index 00000000000..533aab34a08 --- /dev/null +++ b/.changeset/queue-combined-concurrency-stats.md @@ -0,0 +1,5 @@ +--- +"@trigger.dev/core": patch +--- + +Queue retrieve and list API responses now report combined concurrency usage. When a queue has a `combinedConcurrencyLimit`, `concurrency.combined` includes the effective cap, the declared base, any active override, and how many runs are in flight across all concurrency keys. diff --git a/.changeset/queue-total-concurrency-limit.md b/.changeset/queue-total-concurrency-limit.md deleted file mode 100644 index 742f4925744..00000000000 --- a/.changeset/queue-total-concurrency-limit.md +++ /dev/null @@ -1,18 +0,0 @@ ---- -"@trigger.dev/sdk": patch -"@trigger.dev/core": patch ---- - -Cap a queue's total concurrency across all of its `concurrencyKey` values with the new `totalConcurrencyLimit` queue option. On a keyed queue, `concurrencyLimit` applies to each key value independently, so ten active keys with a limit of 5 can run 50 at once. `totalConcurrencyLimit` bounds the whole queue while each key still gets at most `concurrencyLimit`. - -```ts -import { queue } from "@trigger.dev/sdk"; - -export const perUserQueue = queue({ - name: "per-user-queue", - concurrencyLimit: 1, - totalConcurrencyLimit: 10, -}); -``` - -Enforcement happens server-side and only applies to runs triggered with a `concurrencyKey`. Servers that have not enabled total concurrency limits accept the option but do not enforce it yet. diff --git a/.changeset/task-concurrency.md b/.changeset/task-concurrency.md new file mode 100644 index 00000000000..32ecd9b9c50 --- /dev/null +++ b/.changeset/task-concurrency.md @@ -0,0 +1,20 @@ +--- +"@trigger.dev/sdk": patch +"@trigger.dev/core": patch +--- + +Control a task's concurrency with the new `concurrency` option, and share limits across tasks with named concurrency limits. An inline shape caps the task itself; `concurrencyLimit()` declares a limit any task can hold (up to two named limits per task), and a trigger call can switch a run's named limits with its own `concurrency` option. + +```ts +import { concurrencyLimit, task } from "@trigger.dev/sdk"; + +export const openaiLimit = concurrencyLimit({ name: "openai", total: 25 }); + +export const generateSummary = task({ + id: "generate-summary", + concurrency: [{ perKey: 1, total: 5 }, openaiLimit], + run: async (payload) => {}, +}); +``` + +`perKey` caps each `concurrencyKey` pool and `total` caps across everything, keys or not. The queue-level `concurrencyLimit` option keeps working unchanged and is deprecated in favor of `concurrency`. Enforcement happens server-side; servers without support accept the option but do not enforce it yet. diff --git a/apps/webapp/app/runEngine/concerns/queues.server.ts b/apps/webapp/app/runEngine/concerns/queues.server.ts index 1374a34d288..9214e91b9b3 100644 --- a/apps/webapp/app/runEngine/concerns/queues.server.ts +++ b/apps/webapp/app/runEngine/concerns/queues.server.ts @@ -23,7 +23,12 @@ import { Namespace, } from "@internal/cache"; import { singleton } from "~/utils/singleton"; -import type { TaskMetadataCache, TaskMetadataEntry } from "~/services/taskMetadataCache.server"; +import { + parseTaskGates, + type TaskMetadataCache, + type TaskMetadataEntry, + type TaskMetadataGate, +} from "~/services/taskMetadataCache.server"; import { taskMetadataCacheInstance } from "~/services/taskMetadataCacheInstance.server"; import { recordTaskMetaResolve, @@ -95,6 +100,7 @@ export class DefaultQueueManager implements QueueManager { let lockedQueueId: string | undefined; let taskTtl: string | null | undefined; let taskKind: string | undefined; + let taskGates: TaskMetadataGate[] | null | undefined; // Determine queue name based on lockToVersion and provided options if (lockedBackgroundWorker) { @@ -146,6 +152,7 @@ export class DefaultQueueManager implements QueueManager { taskTtl = lockedMeta?.ttl ?? undefined; } taskKind = lockedMeta?.triggerSource; + taskGates = lockedMeta?.gates; } else { // No queue override - resolve default queue + TTL + triggerSource via cache, // falling back to a single BackgroundWorkerTask lookup on miss. @@ -184,6 +191,7 @@ export class DefaultQueueManager implements QueueManager { queueName = lockedMeta.queueName; lockedQueueId = lockedMeta.queueId ?? undefined; taskKind = lockedMeta.triggerSource; + taskGates = lockedMeta.gates; } } else { // Task is not locked to a specific version, use regular logic @@ -199,6 +207,7 @@ export class DefaultQueueManager implements QueueManager { queueName = taskInfo.queueName; taskTtl = taskInfo.taskTtl; taskKind = taskInfo.taskKind; + taskGates = taskInfo.taskGates; } // Sanitize the final determined queue name once @@ -211,17 +220,29 @@ export class DefaultQueueManager implements QueueManager { queueName = sanitizedQueueName; } + const requestedGates = request.body.options?.gates ?? taskGates ?? undefined; + const gates = requestedGates + ?.flatMap((gate) => { + const sanitized = sanitizeQueueName(gate.queue); + return sanitized ? [{ queue: sanitized, concurrencyKey: gate.concurrencyKey }] : []; + }) + .slice(0, 2); + return { queueName, lockedQueueId, taskTtl, taskKind, + gates: gates && gates.length > 0 ? gates : undefined, }; } - private async getTaskQueueInfo( - request: TriggerTaskRequest - ): Promise<{ queueName: string; taskTtl?: string | null; taskKind?: string | undefined }> { + private async getTaskQueueInfo(request: TriggerTaskRequest): Promise<{ + queueName: string; + taskTtl?: string | null; + taskKind?: string | undefined; + taskGates?: TaskMetadataGate[] | null; + }> { const { taskId, environment, body } = request; const { queue } = body.options ?? {}; @@ -243,6 +264,7 @@ export class DefaultQueueManager implements QueueManager { queueName: overriddenQueueName, taskTtl: meta?.ttl ?? undefined, taskKind: meta?.triggerSource, + taskGates: meta?.gates, }; } @@ -259,10 +281,20 @@ export class DefaultQueueManager implements QueueManager { taskId, environmentId: environment.id, }); - return { queueName: defaultQueueName, taskTtl: meta.ttl, taskKind: meta.triggerSource }; + return { + queueName: defaultQueueName, + taskTtl: meta.ttl, + taskKind: meta.triggerSource, + taskGates: meta.gates, + }; } - return { queueName: meta.queueName, taskTtl: meta.ttl, taskKind: meta.triggerSource }; + return { + queueName: meta.queueName, + taskTtl: meta.ttl, + taskKind: meta.triggerSource, + taskGates: meta.gates, + }; } /** @@ -320,6 +352,7 @@ export class DefaultQueueManager implements QueueManager { triggerSource: row.triggerSource, queueId: row.queue?.id ?? null, queueName: row.queue?.name ?? "", + gates: parseTaskGates(row.gates), }; // Fire-and-forget back-fill — `setByWorker` upserts the single field and @@ -340,6 +373,7 @@ export class DefaultQueueManager implements QueueManager { select: { ttl: true, triggerSource: true, + gates: true, queue: { select: { id: true, name: true } }, }, }); @@ -378,6 +412,7 @@ export class DefaultQueueManager implements QueueManager { select: { ttl: true, triggerSource: true, + gates: true, queue: { select: { id: true, name: true } }, }, }); @@ -395,6 +430,7 @@ export class DefaultQueueManager implements QueueManager { triggerSource: row.triggerSource, queueId: row.queue?.id ?? null, queueName: row.queue?.name ?? "", + gates: parseTaskGates(row.gates), }; // Fire-and-forget back-fill — atomically upserts the slug into both diff --git a/apps/webapp/app/runEngine/services/triggerTask.server.ts b/apps/webapp/app/runEngine/services/triggerTask.server.ts index d3320dbc219..2994f3ce7af 100644 --- a/apps/webapp/app/runEngine/services/triggerTask.server.ts +++ b/apps/webapp/app/runEngine/services/triggerTask.server.ts @@ -445,7 +445,7 @@ export class RunEngineTriggerTaskService { const parkedOnExternalDeploymentId = externalDeploymentResolution?.outcome === "park" ? externalDeploymentId : undefined; - const { queueName, lockedQueueId, taskTtl, taskKind } = + const { queueName, lockedQueueId, taskTtl, taskKind, gates } = await this.queueConcern.resolveQueueProperties( triggerRequest, lockedToBackgroundWorker ?? undefined @@ -663,6 +663,7 @@ export class RunEngineTriggerTaskService { options, queueName, lockedQueueId, + gates, workerQueue, region: migrated.region, enableFastPath: migrated.enableFastPath, @@ -743,6 +744,7 @@ export class RunEngineTriggerTaskService { options, queueName, lockedQueueId, + gates, workerQueue, region: migrated.region, enableFastPath: migrated.enableFastPath, @@ -905,6 +907,7 @@ export class RunEngineTriggerTaskService { options: TriggerTaskServiceOptions; queueName: string; lockedQueueId?: string; + gates?: Array<{ queue: string; concurrencyKey?: string }>; workerQueue?: string; region?: string; enableFastPath: boolean; @@ -971,6 +974,7 @@ export class RunEngineTriggerTaskService { : args.body.options?.concurrencyKey, queue: args.queueName, lockedQueueId: args.lockedQueueId, + gates: args.gates, workerQueue: args.workerQueue, region: args.region, enableFastPath: args.enableFastPath, diff --git a/apps/webapp/app/runEngine/types.ts b/apps/webapp/app/runEngine/types.ts index 4e415483120..a30606e92a6 100644 --- a/apps/webapp/app/runEngine/types.ts +++ b/apps/webapp/app/runEngine/types.ts @@ -46,6 +46,8 @@ export type QueueProperties = { lockedQueueId?: string; taskTtl?: string | null; taskKind?: string; + /** Other queues the run must also hold a concurrency slot in while executing. */ + gates?: Array<{ queue: string; concurrencyKey?: string }>; }; export type LockedBackgroundWorker = Pick< diff --git a/apps/webapp/app/services/taskMetadataCache.server.ts b/apps/webapp/app/services/taskMetadataCache.server.ts index 6130295a73f..a9f75f29b9d 100644 --- a/apps/webapp/app/services/taskMetadataCache.server.ts +++ b/apps/webapp/app/services/taskMetadataCache.server.ts @@ -1,13 +1,18 @@ import type { Redis, Result, Callback } from "ioredis"; +import { parseGates } from "@internal/run-engine"; import type { TaskTriggerSource } from "@trigger.dev/database"; import { logger } from "./logger.server"; +export type TaskMetadataGate = { queue: string; concurrencyKey?: string }; + export type TaskMetadataEntry = { slug: string; ttl: string | null; triggerSource: TaskTriggerSource; queueId: string | null; queueName: string; + /** Task-declared gates, applied to every trigger that does not override them. */ + gates: TaskMetadataGate[] | null; }; export interface TaskMetadataCache { @@ -52,11 +57,21 @@ export type RedisTaskMetadataCacheOptions = { byWorkerTtlSeconds?: number; }; +/** + * BackgroundWorkerTask.gates is an untyped Json column; keep only well-shaped + * entries so a malformed value can never fail a trigger. + */ +export function parseTaskGates(gates: unknown): TaskMetadataGate[] | null { + const parsed = parseGates(gates); + return parsed.length > 0 ? parsed : null; +} + type EncodedEntry = { t: string | null; k: TaskTriggerSource; q: string | null; n: string; + g?: TaskMetadataGate[] | null; }; function encode(entry: TaskMetadataEntry): string { @@ -65,6 +80,7 @@ function encode(entry: TaskMetadataEntry): string { k: entry.triggerSource, q: entry.queueId, n: entry.queueName, + g: entry.gates, }; return JSON.stringify(payload); } @@ -78,6 +94,7 @@ function decode(slug: string, raw: string): TaskMetadataEntry | null { triggerSource: parsed.k, queueId: parsed.q, queueName: parsed.n, + gates: parseTaskGates(parsed.g ?? null), }; } catch (error) { logger.error("Failed to decode task metadata cache entry", { slug, error }); diff --git a/apps/webapp/app/v3/services/changeCurrentDeployment.server.ts b/apps/webapp/app/v3/services/changeCurrentDeployment.server.ts index 5f5697dd7cb..021cac19174 100644 --- a/apps/webapp/app/v3/services/changeCurrentDeployment.server.ts +++ b/apps/webapp/app/v3/services/changeCurrentDeployment.server.ts @@ -6,6 +6,7 @@ import { logger } from "~/services/logger.server"; import { syncTaskIdentifiers } from "~/services/taskIdentifierRegistry.server"; import { type TaskMetadataCache, + parseTaskGates, type TaskMetadataEntry, } from "~/services/taskMetadataCache.server"; import { taskMetadataCacheInstance } from "~/services/taskMetadataCacheInstance.server"; @@ -119,6 +120,7 @@ export class ChangeCurrentDeploymentService extends BaseService { slug: true, triggerSource: true, ttl: true, + gates: true, queue: { select: { id: true, name: true } }, }, }) @@ -157,6 +159,7 @@ export class ChangeCurrentDeploymentService extends BaseService { triggerSource: t.triggerSource, queueId: t.queue?.id ?? null, queueName: t.queue?.name ?? "", + gates: parseTaskGates(t.gates), })); // Cache calls log+swallow internally. diff --git a/apps/webapp/app/v3/services/createBackgroundWorker.server.ts b/apps/webapp/app/v3/services/createBackgroundWorker.server.ts index f1cebe640d5..386f8f038d6 100644 --- a/apps/webapp/app/v3/services/createBackgroundWorker.server.ts +++ b/apps/webapp/app/v3/services/createBackgroundWorker.server.ts @@ -403,9 +403,9 @@ async function createWorkerTask( { name: task.queue?.name ?? `task/${task.id}`, concurrencyLimit: task.queue?.concurrencyLimit, - totalConcurrencyLimit: task.queue?.totalConcurrencyLimit, + combinedConcurrencyLimit: task.queue?.combinedConcurrencyLimit, }, - task.id, + task.queue?.name ?? task.id, task.queue?.name ? "NAMED" : "VIRTUAL", worker, environment, @@ -437,6 +437,7 @@ async function createWorkerTask( exportName: task.exportName, retryConfig: task.retry, queueConfig: task.queue, + gates: task.gates, machineConfig: task.machine, triggerSource: resolvedTriggerSource, config: task.agentConfig ? (task.agentConfig as any) : undefined, @@ -454,6 +455,7 @@ async function createWorkerTask( triggerSource: resolvedTriggerSource, queueId: queue.id, queueName: queue.name, + gates: task.gates ?? null, }; } catch (error) { if (error instanceof Prisma.PrismaClientKnownRequestError) { @@ -477,6 +479,7 @@ async function createWorkerTask( triggerSource: resolvedTriggerSource, queueId: queue.id, queueName: queue.name, + gates: task.gates ?? null, }; } } else { @@ -555,7 +558,7 @@ async function createWorkerQueue( const taskQueue = await upsertWorkerQueueRecord( queueName, baseConcurrencyLimit ?? null, - queue.totalConcurrencyLimit ?? null, + queue.combinedConcurrencyLimit ?? null, orderableName, queueType, worker, diff --git a/apps/webapp/app/v3/services/replayTaskRun.server.ts b/apps/webapp/app/v3/services/replayTaskRun.server.ts index 750427fb327..e0f3588bcaa 100644 --- a/apps/webapp/app/v3/services/replayTaskRun.server.ts +++ b/apps/webapp/app/v3/services/replayTaskRun.server.ts @@ -142,6 +142,9 @@ export class ReplayTaskRunService extends BaseService { : undefined, concurrencyKey: overrideOptions.concurrencyKey ?? existingTaskRun.concurrencyKey ?? undefined, + gates: Array.isArray(existingTaskRun.gates) + ? (existingTaskRun.gates as Array<{ queue: string; concurrencyKey?: string }>) + : [], maxAttempts: overrideOptions.maxAttempts, maxDuration: overrideOptions.maxDurationSeconds, machine: diff --git a/internal-packages/database/prisma/migrations/20260829120000_add_queue_gates/migration.sql b/internal-packages/database/prisma/migrations/20260829120000_add_queue_gates/migration.sql new file mode 100644 index 00000000000..d4dfce733d2 --- /dev/null +++ b/internal-packages/database/prisma/migrations/20260829120000_add_queue_gates/migration.sql @@ -0,0 +1,5 @@ +-- AlterTable +ALTER TABLE "BackgroundWorkerTask" ADD COLUMN "gates" JSONB; + +-- AlterTable +ALTER TABLE "TaskRun" ADD COLUMN "gates" JSONB; diff --git a/internal-packages/database/prisma/schema.prisma b/internal-packages/database/prisma/schema.prisma index ac0ac09b2f0..c88ba5b5887 100644 --- a/internal-packages/database/prisma/schema.prisma +++ b/internal-packages/database/prisma/schema.prisma @@ -741,6 +741,9 @@ model BackgroundWorkerTask { queueConfig Json? retryConfig Json? machineConfig Json? + /// Gates declared on the task: other queues its runs must also hold a concurrency + /// slot in while executing. Shape: [{ queue: string, concurrencyKey?: string }] + gates Json? queueId String? queue TaskQueue? @relation(fields: [queueId], references: [id], onDelete: SetNull, onUpdate: Cascade) @@ -1107,6 +1110,10 @@ model TaskRun { concurrencyKey String? + /// Gates for this run: other queues it must also hold a concurrency slot in while + /// executing. Shape: [{ queue: string, concurrencyKey?: string }] + gates Json? + delayUntil DateTime? queuedAt DateTime? ttl String? diff --git a/internal-packages/run-engine/src/engine/gateParsing.test.ts b/internal-packages/run-engine/src/engine/gateParsing.test.ts new file mode 100644 index 00000000000..495b7b5b9f0 --- /dev/null +++ b/internal-packages/run-engine/src/engine/gateParsing.test.ts @@ -0,0 +1,59 @@ +import { describe, expect, it } from "vitest"; +import { parseGates } from "./gateParsing.js"; + +describe("parseGates", () => { + it("keeps well-shaped gates and caps at three", () => { + expect( + parseGates([ + { queue: "a" }, + { queue: "b", concurrencyKey: "shared" }, + { queue: "c" }, + { queue: "d" }, + ]) + ).toEqual([ + { queue: "a", concurrencyKey: undefined }, + { queue: "b", concurrencyKey: "shared" }, + { queue: "c", concurrencyKey: undefined }, + ]); + }); + + it("returns empty for non-arrays and empty arrays", () => { + expect(parseGates(undefined)).toEqual([]); + expect(parseGates(null)).toEqual([]); + expect(parseGates("gates")).toEqual([]); + expect(parseGates([])).toEqual([]); + }); + + it("drops malformed entries", () => { + expect(parseGates([null, "x", 4, { concurrencyKey: "k" }, { queue: 7 }])).toEqual([]); + }); + + it("drops empty and over-length queue names, keeping exactly 128", () => { + const max = "q".repeat(128); + expect(parseGates([{ queue: "" }, { queue: "q".repeat(129) }, { queue: max }])).toEqual([ + { queue: max, concurrencyKey: undefined }, + ]); + }); + + it("treats an empty-string key as omitted and keeps exactly 128-char keys", () => { + const maxKey = "k".repeat(128); + expect(parseGates([{ queue: "a", concurrencyKey: "" }])).toEqual([ + { queue: "a", concurrencyKey: undefined }, + ]); + expect(parseGates([{ queue: "a", concurrencyKey: maxKey }])).toEqual([ + { queue: "a", concurrencyKey: maxKey }, + ]); + }); + + it("drops gates whose literal key exceeds the cap", () => { + expect(parseGates([{ queue: "a", concurrencyKey: "k".repeat(129) }, { queue: "b" }])).toEqual([ + { queue: "b", concurrencyKey: undefined }, + ]); + }); + + it("ignores non-string keys", () => { + expect(parseGates([{ queue: "a", concurrencyKey: 5 }])).toEqual([ + { queue: "a", concurrencyKey: undefined }, + ]); + }); +}); diff --git a/internal-packages/run-engine/src/engine/gateParsing.ts b/internal-packages/run-engine/src/engine/gateParsing.ts new file mode 100644 index 00000000000..18ab4d5d8cb --- /dev/null +++ b/internal-packages/run-engine/src/engine/gateParsing.ts @@ -0,0 +1,34 @@ +/** + * The gate contract for stored/untyped gate values (BackgroundWorkerTask.gates, + * TaskRun.gates are Json columns): keep only well-shaped entries so a malformed + * value can never fail a trigger or an enqueue. A gate needs a queue name within + * the manifest bounds (1-128 chars); a literal concurrency key must fit the same + * bounds, and an empty-string key means "omitted" so the gate inherits the run's + * key. At most three gates apply: a task's anonymous inline-limit gate plus two + * named limits. + */ +export type ParsedGate = { queue: string; concurrencyKey?: string }; + +export function parseGates(gates: unknown): ParsedGate[] { + if (!Array.isArray(gates) || gates.length === 0) { + return []; + } + + const parsed = gates.flatMap((gate): ParsedGate[] => { + if (!gate || typeof gate !== "object" || typeof (gate as any).queue !== "string") { + return []; + } + const queue = (gate as any).queue; + if (queue.length === 0 || queue.length > 128) { + return []; + } + const rawKey = (gate as any).concurrencyKey; + if (typeof rawKey === "string" && rawKey.length > 128) { + return []; + } + const concurrencyKey = typeof rawKey === "string" && rawKey.length > 0 ? rawKey : undefined; + return [{ queue, concurrencyKey }]; + }); + + return parsed.slice(0, 3); +} diff --git a/internal-packages/run-engine/src/engine/index.ts b/internal-packages/run-engine/src/engine/index.ts index 0aef09fe982..40e3d7bc336 100644 --- a/internal-packages/run-engine/src/engine/index.ts +++ b/internal-packages/run-engine/src/engine/index.ts @@ -841,6 +841,7 @@ export class RunEngine { sdkVersion, cliVersion, concurrencyKey, + gates, workerQueue, region, enableFastPath, @@ -1019,6 +1020,7 @@ export class RunEngine { sdkVersion, cliVersion, concurrencyKey, + gates, queue, lockedQueueId, workerQueue, diff --git a/internal-packages/run-engine/src/engine/systems/enqueueSystem.ts b/internal-packages/run-engine/src/engine/systems/enqueueSystem.ts index 38c681c511e..ab2f8828384 100644 --- a/internal-packages/run-engine/src/engine/systems/enqueueSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/enqueueSystem.ts @@ -8,6 +8,7 @@ import type { RunStore } from "@internal/run-store"; import { parseNaturalLanguageDuration } from "@trigger.dev/core/v3/isomorphic"; import type { MinimalAuthenticatedEnvironment } from "../../shared/index.js"; import { QUEUED_SNAPSHOT_DESCRIPTION, QUEUED_SNAPSHOT_STATUS } from "../consts.js"; +import { parseGates } from "../gateParsing.js"; import type { ExecutionSnapshotSystem } from "./executionSnapshotSystem.js"; import type { SystemResources } from "./systems.js"; @@ -16,6 +17,13 @@ export type EnqueueSystemOptions = { executionSnapshotSystem: ExecutionSnapshotSystem; }; +function parseRunGates( + gates: unknown +): Array<{ queue: string; concurrencyKey?: string }> | undefined { + const parsed = parseGates(gates); + return parsed.length > 0 ? parsed : undefined; +} + export class EnqueueSystem { private readonly $: SystemResources; private readonly executionSnapshotSystem: ExecutionSnapshotSystem; @@ -177,6 +185,7 @@ export class EnqueueSystem { environmentType: env.type, queue: run.queue, concurrencyKey: run.concurrencyKey ?? undefined, + gates: parseRunGates(run.gates), timestamp, eligibleAtMs, attempt: 0, diff --git a/internal-packages/run-engine/src/engine/tests/queueGates.test.ts b/internal-packages/run-engine/src/engine/tests/queueGates.test.ts new file mode 100644 index 00000000000..8c06e9d8563 --- /dev/null +++ b/internal-packages/run-engine/src/engine/tests/queueGates.test.ts @@ -0,0 +1,137 @@ +import { containerTest } from "@internal/testcontainers"; +import { trace } from "@internal/tracing"; +import { setTimeout } from "timers/promises"; +import { expect } from "vitest"; +import { RunEngine } from "../index.js"; +import { setupAuthenticatedEnvironment, setupBackgroundWorker } from "./setup.js"; + +vi.setConfig({ testTimeout: 60_000 }); + +async function waitFor(condition: () => Promise, timeoutMs = 20_000): Promise { + const deadline = Date.now() + timeoutMs; + while (Date.now() < deadline) { + if (await condition()) { + return true; + } + await setTimeout(250); + } + return condition(); +} + +describe("RunEngine queue gates", () => { + containerTest( + "trigger persists gates and the queue enforces them", + async ({ prisma, redisOptions }) => { + const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + + const engine = new RunEngine({ + prisma, + worker: { + redis: redisOptions, + workers: 1, + tasksPerWorker: 10, + pollIntervalMs: 100, + }, + queue: { + redis: redisOptions, + processWorkerQueueDebounceMs: 50, + gatesEnabled: true, + }, + runLock: { + redis: redisOptions, + }, + machines: { + defaultMachine: "small-1x", + machines: { + "small-1x": { + name: "small-1x" as const, + cpu: 0.5, + memory: 0.5, + centsPerMs: 0.0001, + }, + }, + baseCostInCents: 0.0001, + }, + tracer: trace.getTracer("test", "0.0.0"), + }); + + try { + const taskIdentifier = "test-task"; + + await setupBackgroundWorker(engine, authenticatedEnvironment, taskIdentifier); + + await engine.runQueue.updateQueueConcurrencyLimits( + authenticatedEnvironment, + "shared-gate", + 1 + ); + + const run1 = await engine.trigger( + { + number: 1, + friendlyId: "run_g1", + environment: authenticatedEnvironment, + taskIdentifier, + payload: "{}", + payloadType: "application/json", + context: {}, + traceContext: {}, + traceId: "t1", + spanId: "s1", + workerQueue: "main", + queue: `task/${taskIdentifier}`, + gates: [{ queue: "shared-gate" }], + isTest: false, + tags: [], + }, + prisma + ); + + const run2 = await engine.trigger( + { + number: 2, + friendlyId: "run_g2", + environment: authenticatedEnvironment, + taskIdentifier, + payload: "{}", + payloadType: "application/json", + context: {}, + traceContext: {}, + traceId: "t2", + spanId: "s2", + workerQueue: "main", + queue: `task/${taskIdentifier}`, + gates: [{ queue: "shared-gate" }], + isTest: false, + tags: [], + }, + prisma + ); + + const storedRun = await prisma.taskRun.findFirst({ where: { id: run1.id } }); + expect(storedRun?.gates).toEqual([{ queue: "shared-gate" }]); + + const oneAdmitted = await waitFor( + async () => + (await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + "shared-gate" + )) === 1 + ); + expect(oneAdmitted).toBe(true); + + /** The second run must stay queued behind the full gate. */ + await setTimeout(2000); + expect( + await engine.runQueue.currentConcurrencyOfQueue(authenticatedEnvironment, "shared-gate") + ).toBe(1); + expect( + await engine.runQueue.lengthOfQueue(authenticatedEnvironment, `task/${taskIdentifier}`) + ).toBe(1); + expect(run2.id).toBeDefined(); + } finally { + await engine.quit(); + } + } + ); +}); diff --git a/internal-packages/run-engine/src/engine/types.ts b/internal-packages/run-engine/src/engine/types.ts index 028858909d9..15dc57cd378 100644 --- a/internal-packages/run-engine/src/engine/types.ts +++ b/internal-packages/run-engine/src/engine/types.ts @@ -23,6 +23,7 @@ import type { workerCatalog } from "./workerCatalog.js"; import { type BillingPlan } from "./billingCache.js"; import type { DRRConfig } from "../batch-queue/types.js"; import type { PendingVersionRunIdLookup } from "./services/pendingVersionLookup.js"; +import type { QueueGate } from "../run-queue/types.js"; /** * Structural mirror of the webapp's CrossSeamGuardDecision @@ -331,6 +332,8 @@ export type TriggerParams = { sdkVersion?: string; cliVersion?: string; concurrencyKey?: string; + /** Other queues this run must also hold a concurrency slot in while executing. At most two. */ + gates?: QueueGate[]; workerQueue?: string; region?: string; /** When true, the run queue may push directly to the worker queue if concurrency is available. diff --git a/internal-packages/run-engine/src/index.ts b/internal-packages/run-engine/src/index.ts index 2c98edf6866..f30e271d5a4 100644 --- a/internal-packages/run-engine/src/index.ts +++ b/internal-packages/run-engine/src/index.ts @@ -61,3 +61,4 @@ export type { WatcherEntry, } from "./engine/waitpointCoordinator/storeCoordinator.js"; export { WaitpointKeyTagError } from "./engine/waitpointCoordinator/keys.js"; +export { parseGates, type ParsedGate } from "./engine/gateParsing.js"; diff --git a/internal-packages/run-engine/src/run-queue/types.ts b/internal-packages/run-engine/src/run-queue/types.ts index ef4a0c2b775..df7cc50b221 100644 --- a/internal-packages/run-engine/src/run-queue/types.ts +++ b/internal-packages/run-engine/src/run-queue/types.ts @@ -9,11 +9,11 @@ import type { MinimalAuthenticatedEnvironment } from "../shared/index.js"; * the entry is keyed) and an extra slot held until release. `queue` is the bare * queue name; the org/project/env scope comes from the run's own payload. */ -const QueueGate = z.object({ +export const QueueGate = z.object({ queue: z.string().min(1).max(128), concurrencyKey: z.string().min(1).max(128).optional(), }); -type QueueGate = z.infer; +export type QueueGate = z.infer; export const InputPayload = z.object({ runId: z.string(), diff --git a/internal-packages/run-ops-database/prisma/migrations/20260829120000_add_task_run_gates/migration.sql b/internal-packages/run-ops-database/prisma/migrations/20260829120000_add_task_run_gates/migration.sql new file mode 100644 index 00000000000..72f07a937f3 --- /dev/null +++ b/internal-packages/run-ops-database/prisma/migrations/20260829120000_add_task_run_gates/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "TaskRun" ADD COLUMN "gates" JSONB; diff --git a/internal-packages/run-ops-database/prisma/schema.prisma b/internal-packages/run-ops-database/prisma/schema.prisma index 0a6c10dd034..b7a1fc77e4c 100644 --- a/internal-packages/run-ops-database/prisma/schema.prisma +++ b/internal-packages/run-ops-database/prisma/schema.prisma @@ -153,6 +153,10 @@ model TaskRun { concurrencyKey String? + /// Gates for this run: other queues it must also hold a concurrency slot in while + /// executing. Shape: [{ queue: string, concurrencyKey?: string }] + gates Json? + delayUntil DateTime? queuedAt DateTime? ttl String? diff --git a/internal-packages/run-store/src/types.ts b/internal-packages/run-store/src/types.ts index 41fecf90bb5..35670d1b92a 100644 --- a/internal-packages/run-store/src/types.ts +++ b/internal-packages/run-store/src/types.ts @@ -174,6 +174,8 @@ export type CreateRunData = { sdkVersion?: string; cliVersion?: string; concurrencyKey?: string; + /** Other queues this run must also hold a concurrency slot in while executing. */ + gates?: Array<{ queue: string; concurrencyKey?: string }>; queue: string; lockedQueueId?: string; workerQueue?: string; diff --git a/packages/cli-v3/src/entryPoints/dev-index-worker.ts b/packages/cli-v3/src/entryPoints/dev-index-worker.ts index 59228f0971d..e093cb34e3b 100644 --- a/packages/cli-v3/src/entryPoints/dev-index-worker.ts +++ b/packages/cli-v3/src/entryPoints/dev-index-worker.ts @@ -196,6 +196,7 @@ await sendMessageInCatalog( prompts: convertPromptSchemasToJsonSchemas(resourceCatalog.listPromptManifests()), skills: resourceCatalog.listSkillManifests(), queues: resourceCatalog.listQueueManifests(), + concurrencyLimits: resourceCatalog.listConcurrencyLimitManifests(), configPath: buildManifest.configPath, runtime: buildManifest.runtime, runtimeVersion: detectRuntimeVersion(), diff --git a/packages/cli-v3/src/entryPoints/managed-index-worker.ts b/packages/cli-v3/src/entryPoints/managed-index-worker.ts index f463c4156e5..866dd2e2453 100644 --- a/packages/cli-v3/src/entryPoints/managed-index-worker.ts +++ b/packages/cli-v3/src/entryPoints/managed-index-worker.ts @@ -192,6 +192,7 @@ await sendMessageInCatalog( prompts: convertPromptSchemasToJsonSchemas(resourceCatalog.listPromptManifests()), skills: resourceCatalog.listSkillManifests(), queues: resourceCatalog.listQueueManifests(), + concurrencyLimits: resourceCatalog.listConcurrencyLimitManifests(), configPath: buildManifest.configPath, runtime: buildManifest.runtime, runtimeVersion: detectRuntimeVersion(), diff --git a/packages/core/src/v3/resource-catalog/catalog.ts b/packages/core/src/v3/resource-catalog/catalog.ts index 1295b3e9144..79ab506fa3b 100644 --- a/packages/core/src/v3/resource-catalog/catalog.ts +++ b/packages/core/src/v3/resource-catalog/catalog.ts @@ -7,6 +7,7 @@ import type { WebhookManifest, WebhookMetadata, WorkerManifest, + ConcurrencyLimitManifest, } from "../schemas/index.js"; import type { PromptMetadataWithFunctions, @@ -27,6 +28,8 @@ export interface ResourceCatalog { registerWorkerManifest(workerManifest: WorkerManifest): void; registerQueueMetadata(queue: QueueManifest): void; listQueueManifests(): Array; + registerConcurrencyLimitMetadata(limit: ConcurrencyLimitManifest): void; + listConcurrencyLimitManifests(): Array; getTaskSchema(id: string): TaskSchema | undefined; registerPromptMetadata(prompt: PromptMetadataWithFunctions): void; listPromptManifests(): Array; diff --git a/packages/core/src/v3/resource-catalog/index.ts b/packages/core/src/v3/resource-catalog/index.ts index 41211892e2c..961c5ccdcd5 100644 --- a/packages/core/src/v3/resource-catalog/index.ts +++ b/packages/core/src/v3/resource-catalog/index.ts @@ -9,6 +9,7 @@ import type { WebhookManifest, WebhookMetadata, WorkerManifest, + ConcurrencyLimitManifest, } from "../schemas/index.js"; import type { PromptMetadataWithFunctions, @@ -94,6 +95,14 @@ export class ResourceCatalogAPI { return this.#getCatalog().listQueueManifests(); } + public registerConcurrencyLimitMetadata(limit: ConcurrencyLimitManifest): void { + this.#getCatalog().registerConcurrencyLimitMetadata(limit); + } + + public listConcurrencyLimitManifests(): Array { + return this.#getCatalog().listConcurrencyLimitManifests(); + } + public registerPromptMetadata(prompt: PromptMetadataWithFunctions): void { this.#getCatalog().registerPromptMetadata(prompt); } diff --git a/packages/core/src/v3/resource-catalog/noopResourceCatalog.ts b/packages/core/src/v3/resource-catalog/noopResourceCatalog.ts index 5fe4d2144d9..8c764ea26c2 100644 --- a/packages/core/src/v3/resource-catalog/noopResourceCatalog.ts +++ b/packages/core/src/v3/resource-catalog/noopResourceCatalog.ts @@ -7,6 +7,7 @@ import type { WebhookManifest, WebhookMetadata, WorkerManifest, + ConcurrencyLimitManifest, } from "../schemas/index.js"; import { type PromptMetadataWithFunctions, @@ -72,6 +73,12 @@ export class NoopResourceCatalog implements ResourceCatalog { return []; } + registerConcurrencyLimitMetadata(limit: ConcurrencyLimitManifest): void {} + + listConcurrencyLimitManifests(): Array { + return []; + } + registerPromptMetadata(prompt: PromptMetadataWithFunctions): void { // noop } diff --git a/packages/core/src/v3/resource-catalog/standardResourceCatalog.ts b/packages/core/src/v3/resource-catalog/standardResourceCatalog.ts index 16fb665332b..f020fb75be4 100644 --- a/packages/core/src/v3/resource-catalog/standardResourceCatalog.ts +++ b/packages/core/src/v3/resource-catalog/standardResourceCatalog.ts @@ -10,6 +10,7 @@ import type { WebhookMetadata, WorkerManifest, QueueManifest, + ConcurrencyLimitManifest, } from "../schemas/index.js"; import type { PromptMetadataWithFunctions, @@ -41,6 +42,7 @@ export class StandardResourceCatalog implements ResourceCatalog { private _promptSchemas: Map = new Map(); private _currentFileContext?: Omit; private _queueMetadata: Map = new Map(); + private _concurrencyLimitMetadata: Map = new Map(); private _skillMetadata: Map = new Map(); private _skillFileMetadata: Map = new Map(); private _webhookMetadata: Map = new Map(); @@ -71,18 +73,26 @@ export class StandardResourceCatalog implements ResourceCatalog { registerQueueMetadata(queue: QueueManifest): void { const existingQueue = this._queueMetadata.get(queue.name); - //if it exists already AND concurrencyLimit is different, log a warning + //if it exists already with different settings, log a warning and keep the first definition if (existingQueue) { const isConcurrencyLimitDifferent = existingQueue.concurrencyLimit !== queue.concurrencyLimit; + const isCombinedLimitDifferent = + existingQueue.combinedConcurrencyLimit !== queue.combinedConcurrencyLimit; - if (isConcurrencyLimitDifferent) { + if (isConcurrencyLimitDifferent || isCombinedLimitDifferent) { let message = `Queue "${queue.name}" is defined twice, with different settings.`; if (isConcurrencyLimitDifferent) { message += `\n - concurrencyLimit: ${existingQueue.concurrencyLimit} vs ${queue.concurrencyLimit}`; } + if (isCombinedLimitDifferent) { + message += `\n - combinedConcurrencyLimit: ${existingQueue.combinedConcurrencyLimit} vs ${queue.combinedConcurrencyLimit}`; + } message += "\n Keeping the first definition:"; message += `\n - concurrencyLimit: ${existingQueue.concurrencyLimit}`; + if (existingQueue.combinedConcurrencyLimit != null) { + message += `\n - combinedConcurrencyLimit: ${existingQueue.combinedConcurrencyLimit}`; + } console.warn(message); return; } @@ -91,6 +101,25 @@ export class StandardResourceCatalog implements ResourceCatalog { this._queueMetadata.set(queue.name, queue); } + registerConcurrencyLimitMetadata(limit: ConcurrencyLimitManifest): void { + const existing = this._concurrencyLimitMetadata.get(limit.name); + + //if it exists already with different settings, log a warning and keep the first definition + if (existing) { + if (existing.perKey !== limit.perKey || existing.total !== limit.total) { + console.warn( + `Concurrency limit "${limit.name}" is defined twice, with different settings.` + + `\n - perKey: ${existing.perKey} vs ${limit.perKey}` + + `\n - total: ${existing.total} vs ${limit.total}` + + `\n Keeping the first definition.` + ); + return; + } + } + + this._concurrencyLimitMetadata.set(limit.name, limit); + } + registerWorkerManifest(workerManifest: WorkerManifest): void { for (const task of workerManifest.tasks) { this._taskFileMetadata.set(task.id, { @@ -214,6 +243,10 @@ export class StandardResourceCatalog implements ResourceCatalog { return Array.from(this._queueMetadata.values()); } + listConcurrencyLimitManifests(): Array { + return Array.from(this._concurrencyLimitMetadata.values()); + } + getTaskManifest(id: string): TaskManifest | undefined { const metadata = this._taskMetadata.get(id); const fileMetadata = this._taskFileMetadata.get(id); diff --git a/packages/core/src/v3/schemas/api.ts b/packages/core/src/v3/schemas/api.ts index 7c539332d1d..b96356b62c3 100644 --- a/packages/core/src/v3/schemas/api.ts +++ b/packages/core/src/v3/schemas/api.ts @@ -327,6 +327,16 @@ export const TriggerTaskRequestBody = z concurrencyLimit: z.number().int().optional(), }) .optional(), + gates: z + .array( + z.object({ + queue: z.string().min(1).max(128), + concurrencyKey: z.string().min(1).max(128).optional(), + }) + ) + .max(3) + .optional(), + concurrency: z.string().min(1).max(128).array().max(2).optional(), concurrencyKey: ConcurrencyKeySchema.optional(), delay: z.string().or(z.coerce.date()).optional(), idempotencyKey: z @@ -433,6 +443,16 @@ export const BatchTriggerTaskItem = z.object({ name: z.string(), }) .optional(), + gates: z + .array( + z.object({ + queue: z.string().min(1).max(128), + concurrencyKey: z.string().min(1).max(128).optional(), + }) + ) + .max(3) + .optional(), + concurrency: z.string().min(1).max(128).array().max(2).optional(), tags: RunTags.optional(), test: z.boolean().optional(), ttl: z.string().or(z.number().nonnegative().int()).optional(), diff --git a/packages/core/src/v3/schemas/build.ts b/packages/core/src/v3/schemas/build.ts index dd4b019003d..80b31e92620 100644 --- a/packages/core/src/v3/schemas/build.ts +++ b/packages/core/src/v3/schemas/build.ts @@ -7,6 +7,7 @@ import { TaskFile, TaskManifest, WebhookManifest, + ConcurrencyLimitManifest, } from "./schemas.js"; export const BuildExternal = z.object({ @@ -115,6 +116,7 @@ export const WorkerManifest = z.object({ webhooks: WebhookManifest.array().optional(), // NEW unclaimedSessionWebhooks: z.array(z.string()).optional(), // session.webhook descriptors no agent listed queues: QueueManifest.array().optional(), + concurrencyLimits: ConcurrencyLimitManifest.array().optional(), workerEntryPoint: z.string(), controllerEntryPoint: z.string().optional(), loaderEntryPoint: z.string().optional(), diff --git a/packages/core/src/v3/schemas/resources.ts b/packages/core/src/v3/schemas/resources.ts index bc0b7a74d8c..84ad36f1071 100644 --- a/packages/core/src/v3/schemas/resources.ts +++ b/packages/core/src/v3/schemas/resources.ts @@ -1,5 +1,12 @@ import { z } from "zod"; -import { QueueManifest, RetryOptions, ScheduleMetadata } from "./schemas.js"; +import { + QueueGateManifest, + QueueManifest, + RetryOptions, + ScheduleMetadata, + TaskConcurrencyManifest, + ConcurrencyLimitManifest, +} from "./schemas.js"; import { MachineConfig } from "./common.js"; import { WebhookVerifierArtifact, @@ -19,6 +26,8 @@ export const TaskResource = z.object({ filePath: z.string(), exportName: z.string().optional(), queue: QueueManifest.extend({ name: z.string().optional() }).optional(), + gates: QueueGateManifest.array().max(2).optional(), + concurrency: TaskConcurrencyManifest.optional(), retry: RetryOptions.optional(), machine: MachineConfig.optional(), triggerSource: z.string().optional(), @@ -80,6 +89,7 @@ export const BackgroundWorkerMetadata = z.object({ prompts: z.array(PromptResource).optional(), webhooks: z.array(WebhookResource).optional(), // NEW queues: z.array(QueueManifest).optional(), + concurrencyLimits: z.array(ConcurrencyLimitManifest).optional(), sourceFiles: z.array(BackgroundWorkerSourceFileMetadata).optional(), runtime: z.string().optional(), runtimeVersion: z.string().optional(), diff --git a/packages/core/src/v3/schemas/schemas.ts b/packages/core/src/v3/schemas/schemas.ts index ff669193c42..05bc1b6b959 100644 --- a/packages/core/src/v3/schemas/schemas.ts +++ b/packages/core/src/v3/schemas/schemas.ts @@ -183,11 +183,49 @@ export const QueueManifest = z.object({ * applies per key value; this is the ceiling for the whole queue. * * Only enforced for runs triggered with a `concurrencyKey`, and requires server-side support. */ - totalConcurrencyLimit: z.number().int().min(0).max(100000).optional().nullable(), + combinedConcurrencyLimit: z.number().int().min(0).max(100000).optional().nullable(), }); export type QueueManifest = z.infer; +/** A gate a task's runs must also hold a concurrency slot in while executing. + * `queue` is the gate queue's name. When `concurrencyKey` is omitted the run's own + * `concurrencyKey` is used, so the gate is keyed per tenant; a literal value pins + * the gate to one shared slot pool. */ +export const QueueGateManifest = z.object({ + queue: z.string().min(1).max(128), + concurrencyKey: z.string().min(1).max(128).optional(), +}); + +export type QueueGateManifest = z.infer; + +/** One limit shape everywhere a limit appears: perKey caps each concurrencyKey pool + * (runs without a key share one pool); total caps across everything, keys or not. + * Limits start at 1 — "block everything" is a pause or an override, not a declared + * limit — so a stored zero always means "no limit" in metrics and engine keys. */ +export const ConcurrencyShapeManifest = z.object({ + perKey: z.number().int().min(1).max(100000).optional(), + total: z.number().int().min(1).max(100000).optional(), +}); + +export type ConcurrencyShapeManifest = z.infer; + +/** A named, shareable concurrency limit declared with concurrencyLimit(). */ +export const ConcurrencyLimitManifest = ConcurrencyShapeManifest.extend({ + name: z.string().min(1).max(128), +}); + +export type ConcurrencyLimitManifest = z.infer; + +/** A task's concurrency declaration: an optional inline limit scoped to the task, + * plus up to two named limits the task's runs also hold while executing. */ +export const TaskConcurrencyManifest = z.object({ + inline: ConcurrencyShapeManifest.optional(), + limits: z.string().min(1).max(128).array().max(2).optional(), +}); + +export type TaskConcurrencyManifest = z.infer; + export const ScheduleMetadata = z.object({ cron: z.string(), timezone: z.string(), @@ -203,6 +241,8 @@ const taskMetadata = { id: z.string(), description: z.string().optional(), queue: QueueManifest.extend({ name: z.string().optional() }).optional(), + gates: QueueGateManifest.array().max(2).optional(), + concurrency: TaskConcurrencyManifest.optional(), retry: RetryOptions.optional(), machine: MachineConfig.optional(), triggerSource: z.string().optional(), diff --git a/packages/core/src/v3/types/queues.ts b/packages/core/src/v3/types/queues.ts index 1fba786c868..fe5e77f94e5 100644 --- a/packages/core/src/v3/types/queues.ts +++ b/packages/core/src/v3/types/queues.ts @@ -31,32 +31,50 @@ export type QueueOptions = { * ``` */ name: string; - /** An optional property that specifies the maximum number of concurrent run executions. - * - * If this property is omitted, the task can potentially use up the full concurrency of an environment */ - concurrencyLimit?: number; - /** An optional property that caps the total number of concurrent run executions across ALL - * `concurrencyKey` values of this queue. - * - * On a queue used with a `concurrencyKey`, `concurrencyLimit` applies to each key value - * independently — ten active keys with `concurrencyLimit: 5` can run 50 at once. Setting - * `totalConcurrencyLimit: 20` bounds the whole queue to 20 while each key still gets at - * most `concurrencyLimit`. - * - * @example - * - * ```ts - * const perUserQueue = queue({ - name: "per-user-queue", - concurrencyLimit: 1, - totalConcurrencyLimit: 10, - }); - * ``` - * - * Only enforced for runs triggered with a `concurrencyKey`, and requires server-side support. - * - * Omit for no total cap. Like `concurrencyLimit`, a value of `0` holds every keyed run in - * the queue rather than removing the cap. + /** + * @deprecated Use `concurrency` on the task instead. `concurrencyLimit: 10` applies per + * `concurrencyKey` when runs pass one, and to the whole queue when they don't. The task's + * `concurrency` option says which you mean: `{ total: 10 }` caps the task outright; + * `{ perKey: 10 }` caps each key. Existing queues keep working unchanged. */ - totalConcurrencyLimit?: number; + concurrencyLimit?: number; }; + +/** + * One limit shape everywhere a limit appears. `perKey` caps each `concurrencyKey` pool + * (runs without a key share one pool); `total` caps across everything, keys or not. + * Either alone or both together. + */ +export type ConcurrencyShape = { + perKey?: number; + total?: number; +}; + +/** Options for `concurrencyLimit()`: a named, shareable concurrency limit. + * + * @example + * + * ```ts + * export const openaiLimit = concurrencyLimit({ name: "openai", total: 25 }); + * + * export const generateSummary = task({ + * id: "generate-summary", + * concurrency: [{ total: 5 }, openaiLimit], + * run: async (payload) => {}, + * }); + * ``` + */ +export type ConcurrencyLimitOptions = { name: string } & ConcurrencyShape; + +export type ConcurrencyLimit = ConcurrencyLimitOptions; + +/** + * A task's concurrency: one limit or an array of limits. An inline shape caps this task; + * a named limit (a `concurrencyLimit()` instance or its name) is shared across every task + * holding it. At most one inline limit plus up to two named limits. + */ +export type TaskConcurrency = + | ConcurrencyShape + | ConcurrencyLimit + | string + | Array; diff --git a/packages/core/src/v3/types/tasks.ts b/packages/core/src/v3/types/tasks.ts index 070cc96f3df..e4ad810a2b4 100644 --- a/packages/core/src/v3/types/tasks.ts +++ b/packages/core/src/v3/types/tasks.ts @@ -26,7 +26,7 @@ import type { TaskRunContext, } from "../schemas/index.js"; import type { IdempotencyKey } from "./idempotencyKeys.js"; -import type { QueueOptions } from "./queues.js"; +import type { QueueOptions, TaskConcurrency } from "./queues.js"; import type { AnySchemaParseFn, inferSchemaIn, inferSchemaOut, Schema } from "./schemas.js"; import type { inferToolParameters, ToolTaskParameters } from "./tools.js"; import type { Prettify } from "./utils.js"; @@ -222,11 +222,25 @@ type CommonTaskOptions< }); * ``` */ - queue?: { - name?: string; - concurrencyLimit?: number; - totalConcurrencyLimit?: number; - }; + queue?: string | TaskQueueIn; + + /** + * Limit how many of this task's runs execute at once, or hold shared named limits. + * Takes one limit or an array: an inline `{ perKey?, total? }` shape caps this task; + * a `concurrencyLimit()` instance (or its name) is shared across every task holding it. + * At most one inline limit plus up to two named limits. + * + * @example + * + * ```ts + * export const sendEmail = task({ + * id: "send-email", + * concurrency: { total: 10 }, + * run: async (payload) => {}, + * }); + * ``` + */ + concurrency?: TaskConcurrency; /** Configure the spec of the [machine](https://trigger.dev/docs/machines) you want your task to run on. * * @example @@ -396,6 +410,22 @@ type CommonTaskOptions< agentConfig?: { type: string }; }; +/** + * A reference to a gate: another queue a run must also hold a concurrency slot in while + * it executes. A plain string names the gate queue; the object form pins the gate to a + * literal `concurrencyKey` instead of inheriting the run's own key. + */ + +type TaskQueueIn = { + name?: string; + /** + * @deprecated Use `concurrency` on the task instead. `concurrencyLimit: 10` applies per + * `concurrencyKey` when runs pass one, and to the whole queue when they don't; the task's + * `concurrency` option says which you mean. Existing queues keep working unchanged. + */ + concurrencyLimit?: number; +}; + export type TaskOptions< TIdentifier extends string, TPayload = void, @@ -812,6 +842,13 @@ export type TriggerOptions = { */ queue?: string; + /** + * Override the task's named concurrency limits for this run. Strings only, like `queue`: + * pass a limit's name (e.g. `paidTier.name`). Replaces the task's declared named limits; + * the task's inline limit always applies. + */ + concurrency?: string | string[]; + /** * The `concurrencyKey` creates a copy of the queue for every unique value of the key. * For example, if the queue (set when triggering or on the task) has a concurrency limit of 10, diff --git a/packages/react-hooks/src/hooks/useTaskTrigger.ts b/packages/react-hooks/src/hooks/useTaskTrigger.ts index 979cd87983b..3f9018596ea 100644 --- a/packages/react-hooks/src/hooks/useTaskTrigger.ts +++ b/packages/react-hooks/src/hooks/useTaskTrigger.ts @@ -73,10 +73,19 @@ export function useTaskTrigger( const payloadPacket = await stringifyIO(payload); + const queueName = options?.queue; + const concurrency = options?.concurrency + ? (Array.isArray(options.concurrency) ? options.concurrency : [options.concurrency]).slice( + 0, + 2 + ) + : undefined; + const handle = await apiClient.triggerTask(id, { payload: payloadPacket.data, options: { - queue: options?.queue ? { name: options.queue } : undefined, + queue: queueName ? { name: queueName } : undefined, + concurrency, concurrencyKey: options?.concurrencyKey, payloadType: payloadPacket.dataType, idempotencyKey: await makeIdempotencyKey(options?.idempotencyKey), diff --git a/packages/trigger-sdk/src/v3/index.ts b/packages/trigger-sdk/src/v3/index.ts index 4bdb582d7ae..c6d940f596d 100644 --- a/packages/trigger-sdk/src/v3/index.ts +++ b/packages/trigger-sdk/src/v3/index.ts @@ -1,7 +1,7 @@ export * from "./cache.js"; export * from "./config.js"; export { retry, type RetryOptions } from "./retry.js"; -export { queue, BatchTriggerError } from "./shared.js"; +export { queue, concurrencyLimit, BatchTriggerError } from "./shared.js"; export * from "./tasks.js"; export * from "./batch.js"; export * from "./wait.js"; diff --git a/packages/trigger-sdk/src/v3/shared.ts b/packages/trigger-sdk/src/v3/shared.ts index b962697d559..1ee4a2e33c1 100644 --- a/packages/trigger-sdk/src/v3/shared.ts +++ b/packages/trigger-sdk/src/v3/shared.ts @@ -90,6 +90,9 @@ import { type TriggerAndWaitOptions, type TriggerApiRequestOptions, type TriggerOptions, + type TaskConcurrency, + type ConcurrencyLimitOptions, + type ConcurrencyLimit, } from "@trigger.dev/core/v3"; import { tracer } from "./tracer.js"; @@ -134,6 +137,105 @@ function resolveTriggerExternalDeploymentId(explicit?: string): string | undefin }); } +type NormalizedTaskQueue = { + queue?: { name?: string; concurrencyLimit?: number }; +}; + +/** A task's `queue` is a queue object or a string reference to a queue defined elsewhere. */ +function normalizeTaskQueue(queue: TaskOptions["queue"]): NormalizedTaskQueue { + if (!queue) { + return {}; + } + + return { queue: typeof queue === "string" ? { name: queue } : queue }; +} + +type NormalizedTaskConcurrency = { + inline?: { perKey?: number; total?: number }; + limits?: string[]; +}; + +/** + * Validates and normalizes the task `concurrency` option into its manifest shape: + * at most one inline limit (caps this task) plus up to two named limits (shared). + */ +function normalizeTaskConcurrency( + taskId: string, + concurrency: TaskConcurrency | undefined +): NormalizedTaskConcurrency | undefined { + if (!concurrency) { + return undefined; + } + + const items = Array.isArray(concurrency) ? concurrency : [concurrency]; + const inline: Array<{ perKey?: number; total?: number }> = []; + const limits: string[] = []; + + for (const item of items) { + if (typeof item === "string") { + limits.push(item); + } else if ( + item && + typeof item === "object" && + "name" in item && + typeof item.name === "string" + ) { + validateConcurrencyLimitName(item.name); + resourceCatalog.registerConcurrencyLimitMetadata(item); + limits.push(item.name); + } else if (item && typeof item === "object") { + inline.push({ perKey: item.perKey, total: item.total }); + } + } + + if (inline.length > 1) { + throw new Error( + `Task "${taskId}": concurrency accepts at most one inline limit. Give shared limits a name with concurrencyLimit().` + ); + } + + if (limits.length > 2) { + throw new Error(`Task "${taskId}": concurrency accepts at most two named limits.`); + } + + return { + inline: inline[0], + limits: limits.length > 0 ? limits : undefined, + }; +} + +/** Builds the queue field of a trigger request body. */ +function triggerQueueBody( + queue: string | undefined, + fallbackQueueName?: string +): { queue?: { name: string } } { + const name = queue ?? fallbackQueueName; + return { queue: name ? { name } : undefined }; +} + +/** + * Trigger-time named limits: strings only, like `queue`. They replace the task's + * declared named limits for this run; the server resolves names to the run's gates. + */ +function triggerConcurrencyBody(concurrency: string | string[] | undefined): { + concurrency?: string[]; +} { + if (!concurrency) { + return {}; + } + const limits = Array.isArray(concurrency) ? concurrency : [concurrency]; + if (limits.length > 2) { + throw new Error("The concurrency option accepts at most two named limits."); + } + if (limits.some((name) => typeof name !== "string" || name.length === 0)) { + throw new Error("The concurrency option takes limit names: non-empty strings."); + } + for (const name of limits) { + validateConcurrencyLimitName(name); + } + return { concurrency: limits }; +} + export function queue(options: QueueOptions): Queue { resourceCatalog.registerQueueMetadata(options); @@ -143,6 +245,41 @@ export function queue(options: QueueOptions): Queue { return options; } +/** + * Declares a named, shareable concurrency limit. Tasks hold it via their `concurrency` + * option; every task holding the same limit draws from the same pools. + * + * @example + * + * ```ts + * export const openaiLimit = concurrencyLimit({ name: "openai", total: 25 }); + * + * export const generateSummary = task({ + * id: "generate-summary", + * concurrency: [{ total: 5 }, openaiLimit], + * run: async (payload) => {}, + * }); + * ``` + */ +function validateConcurrencyLimitName(name: string): void { + if (!/^[a-zA-Z0-9_-]{1,122}$/.test(name)) { + throw new Error( + `Concurrency limit "${name}": names are 1-122 characters using only letters, numbers, underscores and hyphens.` + ); + } +} + +export function concurrencyLimit(options: ConcurrencyLimitOptions): ConcurrencyLimit { + validateConcurrencyLimitName(options.name); + + resourceCatalog.registerConcurrencyLimitMetadata(options); + + // @ts-expect-error + options[Symbol.for("trigger.dev/concurrencyLimit")] = true; + + return options; +} + // Overload: when payloadSchema is provided, payload type should be any export function createTask< TIdentifier extends string, @@ -172,6 +309,8 @@ export function createTask< | TaskOptions | TaskOptionsWithSchema ): Task | Task { + const normalizedQueue = normalizeTaskQueue(params.queue); + const task: Task = { id: params.id, description: params.description, @@ -183,7 +322,7 @@ export function createTask< payload, undefined, { - queue: params.queue?.name, + queue: normalizedQueue.queue?.name, ...options, } ); @@ -196,7 +335,7 @@ export function createTask< options, undefined, undefined, - params.queue?.name + normalizedQueue.queue?.name ); }, triggerAndWait: (payload, options, requestOptions) => { @@ -207,7 +346,7 @@ export function createTask< payload, undefined, { - queue: params.queue?.name, + queue: normalizedQueue.queue?.name, ...options, }, requestOptions @@ -228,7 +367,7 @@ export function createTask< payload, undefined, { - queue: params.queue?.name, + queue: normalizedQueue.queue?.name, ...options, } ) @@ -248,7 +387,7 @@ export function createTask< undefined, options, undefined, - params.queue?.name + normalizedQueue.queue?.name ); }, }; @@ -258,7 +397,8 @@ export function createTask< resourceCatalog.registerTaskMetadata({ id: params.id, description: params.description, - queue: params.queue, + queue: normalizedQueue.queue, + concurrency: normalizeTaskConcurrency(params.id, params.concurrency), retry: params.retry ? { ...defaultRetryOptions, ...params.retry } : undefined, machine: typeof params.machine === "string" ? { preset: params.machine } : params.machine, triggerSource: params.triggerSource, @@ -271,13 +411,19 @@ export function createTask< }, }); - const queue = params.queue; + const queue = normalizedQueue.queue; - if (queue && typeof queue.name === "string") { + /** + * A string queue name (bare or as a tuple's first element) is a REFERENCE to a + * queue defined elsewhere; registering it would create an empty definition that + * can shadow the real one depending on module evaluation order. + */ + const homeIsReference = typeof params.queue === "string"; + + if (queue && typeof queue.name === "string" && !homeIsReference) { resourceCatalog.registerQueueMetadata({ name: queue.name, concurrencyLimit: queue.concurrencyLimit, - totalConcurrencyLimit: queue.totalConcurrencyLimit, }); } @@ -327,6 +473,8 @@ export function createSchemaTask< ? getSchemaParseFn>(params.schema) : undefined; + const normalizedQueue = normalizeTaskQueue(params.queue); + const task: TaskWithSchema = { id: params.id, description: params.description, @@ -338,7 +486,7 @@ export function createSchemaTask< payload, parsePayload, { - queue: params.queue?.name, + queue: normalizedQueue.queue?.name, ...options, }, requestOptions @@ -352,7 +500,7 @@ export function createSchemaTask< options, parsePayload, requestOptions, - params.queue?.name + normalizedQueue.queue?.name ); }, triggerAndWait: (payload, options) => { @@ -363,7 +511,7 @@ export function createSchemaTask< payload, parsePayload, { - queue: params.queue?.name, + queue: normalizedQueue.queue?.name, ...options, } ) @@ -383,7 +531,7 @@ export function createSchemaTask< payload, parsePayload, { - queue: params.queue?.name, + queue: normalizedQueue.queue?.name, ...options, } ) @@ -403,7 +551,7 @@ export function createSchemaTask< parsePayload, options, undefined, - params.queue?.name + normalizedQueue.queue?.name ); }, }; @@ -413,7 +561,8 @@ export function createSchemaTask< resourceCatalog.registerTaskMetadata({ id: params.id, description: params.description, - queue: params.queue, + queue: normalizedQueue.queue, + concurrency: normalizeTaskConcurrency(params.id, params.concurrency), retry: params.retry ? { ...defaultRetryOptions, ...params.retry } : undefined, machine: typeof params.machine === "string" ? { preset: params.machine } : params.machine, triggerSource: params.triggerSource, @@ -427,13 +576,19 @@ export function createSchemaTask< schema: params.schema, }); - const queue = params.queue; + const queue = normalizedQueue.queue; + + /** + * A string queue name (bare or as a tuple's first element) is a REFERENCE to a + * queue defined elsewhere; registering it would create an empty definition that + * can shadow the real one depending on module evaluation order. + */ + const homeIsReference = typeof params.queue === "string"; - if (queue && typeof queue.name === "string") { + if (queue && typeof queue.name === "string" && !homeIsReference) { resourceCatalog.registerQueueMetadata({ name: queue.name, concurrencyLimit: queue.concurrencyLimit, - totalConcurrencyLimit: queue.totalConcurrencyLimit, }); } @@ -731,7 +886,8 @@ export async function batchTriggerById( task: item.id, payload: payloadPacket.data, options: { - queue: item.options?.queue ? { name: item.options.queue } : undefined, + ...triggerQueueBody(item.options?.queue), + ...triggerConcurrencyBody(item.options?.concurrency), concurrencyKey: item.options?.concurrencyKey, test: taskContext.ctx?.run.isTest, payloadType: payloadPacket.dataType, @@ -991,7 +1147,8 @@ export async function batchTriggerByIdAndWait( payload: payloadPacket.data, options: { lockToVersion: taskContext.worker?.version, - queue: item.options?.queue ? { name: item.options.queue } : undefined, + ...triggerQueueBody(item.options?.queue), + ...triggerConcurrencyBody(item.options?.concurrency), concurrencyKey: item.options?.concurrencyKey, test: taskContext.ctx?.run.isTest, payloadType: payloadPacket.dataType, @@ -1256,7 +1413,8 @@ export async function batchTriggerTasks( task: item.task.id, payload: payloadPacket.data, options: { - queue: item.options?.queue ? { name: item.options.queue } : undefined, + ...triggerQueueBody(item.options?.queue), + ...triggerConcurrencyBody(item.options?.concurrency), concurrencyKey: item.options?.concurrencyKey, test: taskContext.ctx?.run.isTest, payloadType: payloadPacket.dataType, @@ -1521,7 +1679,8 @@ export async function batchTriggerAndWaitTasks( task: item.id, payload: payloadPacket.data, options: { - queue: item.options?.queue ? { name: item.options.queue } : undefined, + ...triggerQueueBody(item.options?.queue), + ...triggerConcurrencyBody(item.options?.concurrency), concurrencyKey: item.options?.concurrencyKey, test: taskContext.ctx?.run.isTest, payloadType: payloadPacket.dataType, @@ -2063,7 +2223,8 @@ async function* transformBatchItemsStreamForWait( payload: payloadPacket.data, options: { lockToVersion: taskContext.worker?.version, - queue: item.options?.queue ? { name: item.options.queue } : undefined, + ...triggerQueueBody(item.options?.queue), + ...triggerConcurrencyBody(item.options?.concurrency), concurrencyKey: item.options?.concurrencyKey, test: taskContext.ctx?.run.isTest, payloadType: payloadPacket.dataType, @@ -2113,7 +2274,8 @@ async function* transformBatchByTaskItemsStream( task: taskIdentifier, payload: payloadPacket.data, options: { - queue: item.options?.queue - ? { name: item.options.queue } - : queue - ? { name: queue } - : undefined, + ...triggerQueueBody(item.options?.queue, queue), + ...triggerConcurrencyBody(item.options?.concurrency), concurrencyKey: item.options?.concurrencyKey, test: taskContext.ctx?.run.isTest, payloadType: payloadPacket.dataType, @@ -2280,11 +2440,8 @@ async function* transformSingleTaskBatchItemsStreamForWait( payload: payloadPacket.data, options: { lockToVersion: taskContext.worker?.version, - queue: item.options?.queue - ? { name: item.options.queue } - : queue - ? { name: queue } - : undefined, + ...triggerQueueBody(item.options?.queue, queue), + ...triggerConcurrencyBody(item.options?.concurrency), concurrencyKey: item.options?.concurrencyKey, test: taskContext.ctx?.run.isTest, payloadType: payloadPacket.dataType, @@ -2334,7 +2491,8 @@ async function trigger_internal( { payload: triggerPayloadPacket.data, options: { - queue: options?.queue ? { name: options.queue } : undefined, + ...triggerQueueBody(options?.queue), + ...triggerConcurrencyBody(options?.concurrency), concurrencyKey: options?.concurrencyKey, test: taskContext.ctx?.run.isTest, payloadType: triggerPayloadPacket.dataType, @@ -2419,11 +2577,8 @@ async function batchTrigger_internal( task: taskIdentifier, payload: payloadPacket.data, options: { - queue: item.options?.queue - ? { name: item.options.queue } - : queue - ? { name: queue } - : undefined, + ...triggerQueueBody(item.options?.queue, queue), + ...triggerConcurrencyBody(item.options?.concurrency), concurrencyKey: item.options?.concurrencyKey, test: taskContext.ctx?.run.isTest, payloadType: payloadPacket.dataType, @@ -2603,7 +2758,8 @@ async function triggerAndWait_internal