diff --git a/.changeset/chat-session-concurrency.md b/.changeset/chat-session-concurrency.md new file mode 100644 index 00000000000..4d2f5eca086 --- /dev/null +++ b/.changeset/chat-session-concurrency.md @@ -0,0 +1,12 @@ +--- +"@trigger.dev/sdk": minor +"@trigger.dev/core": minor +--- + +Chat agents can now scope concurrency per session. Pass `concurrencyKey` (for example, your chat ID or tenant ID) and trigger-time named limits via `triggerConfig.concurrency` when starting a chat session, from `chat.createStartSessionAction`, the `AgentChat` client, or a handover. Keys are never defaulted, so a session without one shares the task's keyless pool. + +```ts +const start = chat.createStartSessionAction("support-chat", { + triggerConfig: { concurrencyKey: user.id }, +}); +``` diff --git a/.changeset/steering-history-order.md b/.changeset/steering-history-order.md new file mode 100644 index 00000000000..9c11e6e1cda --- /dev/null +++ b/.changeset/steering-history-order.md @@ -0,0 +1,5 @@ +--- +"@trigger.dev/sdk": patch +--- + +Steering messages now remain in context across agent steps and keep their original position in saved conversations, including custom response data written between steps. diff --git a/.changeset/task-concurrency.md b/.changeset/task-concurrency.md new file mode 100644 index 00000000000..a1ed8bfbe46 --- /dev/null +++ b/.changeset/task-concurrency.md @@ -0,0 +1,25 @@ +--- +"@trigger.dev/sdk": minor +"@trigger.dev/core": minor +"@trigger.dev/react-hooks": minor +--- + +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. + +Manage limits at runtime with the new `concurrencyLimits` namespace: `list()` and `retrieve(name)` report each limit's bounds plus its live `running` and `queued` counts, `override(name, { perKey, total })` changes only the given bounds (overriding `total` to `0` pauses the limit), and `reset(name)` restores the declared values. + +Queue reads (`queues.list()` and `queues.retrieve()`) now report a `version` that discriminates the shape: `V1` queues keep today's fields (their own `concurrencyLimit` and its override state), while `V2` queues (tasks declared with `concurrency`) carry no queue-level concurrency, since their limits are read and overridden through `concurrencyLimits` (a task's inline limit under its derived `task/` name). Existing reads keep compiling: a `V2` queue reports `concurrencyLimit` as null and `concurrency` as undefined. diff --git a/apps/webapp/app/components/billing/OrgBanner.tsx b/apps/webapp/app/components/billing/OrgBanner.tsx index be93fa188c8..eb29f94bdce 100644 --- a/apps/webapp/app/components/billing/OrgBanner.tsx +++ b/apps/webapp/app/components/billing/OrgBanner.tsx @@ -21,10 +21,10 @@ import { useOptionalProject, useProject } from "~/hooks/useProject"; import { useShowSelfServe } from "~/hooks/useShowSelfServe"; import { useCurrentPlan } from "~/routes/_app.orgs.$organizationSlug/route"; import { + concurrencyPath, organizationProjectsPath, v3BillingLimitsPath, v3BillingPath, - v3QueuesPath, } from "~/utils/pathBuilder"; import { ENVIRONMENT_PAUSE_SOURCE_BILLING_LIMIT } from "~/utils/environmentPauseSource"; @@ -71,7 +71,7 @@ export function OrgBanner() { showSelfServe, }); - const hideQueuesButton = location.pathname.endsWith("/queues"); + const hideConcurrencyButton = location.pathname.endsWith("/concurrency"); const hideBillingLimitBanner = location.pathname.endsWith("/settings/billing-limits"); switch (bannerKind) { @@ -89,7 +89,7 @@ export function OrgBanner() { return isArchived ? ( ) : ( - + ); default: return null; @@ -303,7 +303,7 @@ function PausedEnvironmentBanner({ hideButton }: { hideButton: boolean }) { hideButton ? undefined : ( Manage diff --git a/apps/webapp/app/components/dashboard-agent/DashboardAgentPanel.tsx b/apps/webapp/app/components/dashboard-agent/DashboardAgentPanel.tsx index 44e3167a72e..5d8be85d855 100644 --- a/apps/webapp/app/components/dashboard-agent/DashboardAgentPanel.tsx +++ b/apps/webapp/app/components/dashboard-agent/DashboardAgentPanel.tsx @@ -52,7 +52,7 @@ import { } from "./unread-counts"; import { AgentPanelColumn, type DashboardAgentMode, type DragHandleProps } from "./panel-layout"; import { markerAfterActiveChat, markerAfterActivity } from "./thinking-marker"; -import { concurrencyPath } from "~/utils/pathBuilder"; +import { concurrencyLimitsPath } from "~/utils/pathBuilder"; import { scopeMatchesPath, sessionPathFor } from "./agent-scope"; function serializePageContext(pageContext: AgentPageContext): string | undefined { @@ -149,7 +149,7 @@ export function DashboardAgentPanel({ const entityId = agentPageEntityId(pageContext, location.pathname); const pagePaths = useMemo>( - () => ({ raise_env_limit: concurrencyPath(organization, project, environment) }), + () => ({ raise_env_limit: concurrencyLimitsPath(organization, project, environment) }), [organization, project, environment] ); diff --git a/apps/webapp/app/components/dashboard-agent/page-label.ts b/apps/webapp/app/components/dashboard-agent/page-label.ts index d78a881473e..47723072757 100644 --- a/apps/webapp/app/components/dashboard-agent/page-label.ts +++ b/apps/webapp/app/components/dashboard-agent/page-label.ts @@ -24,7 +24,7 @@ const KIND_LABELS: Record, string> = { alerts: "Alerts", apikeys: "API keys", envvars: "Environment variables", - concurrency: "Concurrency", + concurrency: "Concurrency limits", regions: "Regions", settings: "Settings", waitpoints: "Waitpoints", @@ -50,6 +50,7 @@ const SECTION_LABELS: Record = { "bulk-actions": "Bulk actions", branches: "Branches", concurrency: "Concurrency", + "concurrency-limits": "Concurrency limits", dashboards: "Dashboards", deployments: "Deployments", "dev-branches": "Branches", diff --git a/apps/webapp/app/components/navigation/favoritePages.tsx b/apps/webapp/app/components/navigation/favoritePages.tsx index 50668bae889..0b9742dcca2 100644 --- a/apps/webapp/app/components/navigation/favoritePages.tsx +++ b/apps/webapp/app/components/navigation/favoritePages.tsx @@ -223,7 +223,8 @@ const ENV_PAGE_META: Record = { logs: { icon: "logs", name: "Logs" }, errors: { icon: "errors", name: "Errors", singular: "Error" }, query: { icon: "query", name: "Query" }, - queues: { icon: "queues", name: "Queues", singular: "Queue" }, + queues: { icon: "queues", name: "Concurrency" }, + concurrency: { icon: "queues", name: "Concurrency" }, dashboards: { icon: "dashboards", name: "Dashboards", singular: "Dashboard" }, deployments: { icon: "deployments", name: "Deploys", singular: "Deploy" }, "environment-variables": { icon: "environment-variables", name: "Environment variables" }, @@ -234,7 +235,7 @@ const ENV_PAGE_META: Record = { "bulk-actions": { icon: "bulk-actions", name: "Bulk actions", singular: "Bulk action" }, apikeys: { icon: "apikeys", name: "API keys" }, alerts: { icon: "alerts", name: "Alerts", singular: "Alert" }, - concurrency: { icon: "concurrency", name: "Concurrency" }, + "concurrency-limits": { icon: "concurrency", name: "Concurrency limits" }, limits: { icon: "limits", name: "Limits" }, schedules: { icon: "schedules", name: "Schedules", singular: "Schedule" }, test: { icon: "test", name: "Test", singular: "Test" }, diff --git a/apps/webapp/app/components/navigation/sideMenuSections.tsx b/apps/webapp/app/components/navigation/sideMenuSections.tsx index 87d5da61195..6a9c0f7145a 100644 --- a/apps/webapp/app/components/navigation/sideMenuSections.tsx +++ b/apps/webapp/app/components/navigation/sideMenuSections.tsx @@ -23,7 +23,7 @@ import { type OrgForPath, type ProjectForPath, branchesPath, - concurrencyPath, + concurrencyLimitsPath, limitsPath, queryPath, regionsPath, @@ -39,7 +39,7 @@ import { v3ProjectAlertsPath, v3ProjectSettingsIntegrationsPath, v3PromptsPath, - v3QueuesPath, + concurrencyPath, v3WaitpointTokensPath, } from "~/utils/pathBuilder"; import { AlphaBadge, NewBadge } from "../FeatureBadges"; @@ -160,10 +160,10 @@ export function buildSideMenuSections({ } satisfies SideMenuItemConfig, { id: "queues", - name: "Queues", + name: "Concurrency", icon: QueuesIcon, activeIconColor: "text-queues", - to: v3QueuesPath(organization, project, environment), + to: concurrencyPath(organization, project, environment), dataAction: "queues", } satisfies SideMenuItemConfig, { @@ -269,10 +269,10 @@ export function buildSideMenuSections({ ? [ { id: "concurrency", - name: "Concurrency", + name: "Concurrency limits", icon: ConcurrencyIcon, activeIconColor: "text-text-bright", - to: concurrencyPath(organization, project, environment), + to: concurrencyLimitsPath(organization, project, environment), dataAction: "concurrency", } satisfies SideMenuItemConfig, ] diff --git a/apps/webapp/app/components/primitives/charts/MetricChart.tsx b/apps/webapp/app/components/primitives/charts/MetricChart.tsx index 0915f169821..9cfcf48c941 100644 --- a/apps/webapp/app/components/primitives/charts/MetricChart.tsx +++ b/apps/webapp/app/components/primitives/charts/MetricChart.tsx @@ -41,6 +41,8 @@ type MetricChartProps = { * are config values that existed all along, so carry the first value backward instead. */ carryBackfill?: string[]; + /** Column that, when positive on a bucket, exempts it from the carryBackfill overwrite. */ + carryBackfillGuard?: string; /** * Line only. Recolour a series' stroke above a threshold with a gradient split (colour only * above the line). `value` sets a constant threshold; `valueFromSeries` reads a (roughly @@ -78,13 +80,22 @@ export function MetricChart({ timeRange, warningOverlay, carryBackfill, + carryBackfillGuard, thresholdStroke, onHasDataChange, sampleCountColumn, }: MetricChartProps) { const { points, xKind: resolvedXKind } = useMemo( - () => buildMetricPoints(rows, { series, xColumn, xKind, carryBackfill, sampleCountColumn }), - [rows, series, xColumn, xKind, carryBackfill, sampleCountColumn] + () => + buildMetricPoints(rows, { + series, + xColumn, + xKind, + carryBackfill, + carryBackfillGuard, + sampleCountColumn, + }), + [rows, series, xColumn, xKind, carryBackfill, carryBackfillGuard, sampleCountColumn] ); const chartConfig = useMemo(() => { diff --git a/apps/webapp/app/components/primitives/charts/metricPoints.ts b/apps/webapp/app/components/primitives/charts/metricPoints.ts index f87c04db1aa..ac60512660d 100644 --- a/apps/webapp/app/components/primitives/charts/metricPoints.ts +++ b/apps/webapp/app/components/primitives/charts/metricPoints.ts @@ -58,12 +58,22 @@ export type BuildMetricPointsOptions = { /** Defaults to `time` when every x is a datetime or epoch number, `category` otherwise. */ xKind?: MetricXKind; carryBackfill?: string[]; + /** Column that, when positive on a bucket, exempts it from the carryBackfill overwrite: + * the bucket has a real sampled value, so an earlier bucket's carry must not clobber it. */ + carryBackfillGuard?: string; sampleCountColumn?: string; }; export function buildMetricPoints( rows: MetricChartRow[], - { series, xColumn = "t", xKind, carryBackfill, sampleCountColumn }: BuildMetricPointsOptions + { + series, + xColumn = "t", + xKind, + carryBackfill, + carryBackfillGuard, + sampleCountColumn, + }: BuildMetricPointsOptions ): { points: MetricPoint[]; xKind: MetricXKind } { // Rows built by seriesFromRows already carry the coordinate under the reserved key. const xOf = (row: MetricChartRow) => (METRIC_X_KEY in row ? row[METRIC_X_KEY] : row[xColumn]); @@ -80,6 +90,9 @@ export function buildMetricPoints( // A bucket with no value for this series is a gap, not a zero. point[s.key] = hasSamples && value != null ? toNumber(value) : null; } + if (carryBackfillGuard) { + point[carryBackfillGuard] = toNumber(r[carryBackfillGuard]); + } // Set last, so a series key colliding with the reserved key can't displace the coordinate. point[METRIC_X_KEY] = kind === "time" ? (timeValueMs(xOf(r)) ?? NaN) : String(xOf(r) ?? ""); return point; @@ -99,7 +112,10 @@ export function buildMetricPoints( const first = points.findIndex((p) => toNumber(p[key]) > 0); if (first > 0) { const value = points[first]![key]!; - for (let i = 0; i < first; i++) points[i]![key] = value; + for (let i = 0; i < first; i++) { + if (carryBackfillGuard && toNumber(points[i]![carryBackfillGuard]) > 0) continue; + points[i]![key] = value; + } } } } diff --git a/apps/webapp/app/components/queues/QueueControls.tsx b/apps/webapp/app/components/queues/QueueControls.tsx index 00853f62538..2995103c98c 100644 --- a/apps/webapp/app/components/queues/QueueControls.tsx +++ b/apps/webapp/app/components/queues/QueueControls.tsx @@ -1,12 +1,13 @@ import { AdjustmentsHorizontalIcon, PauseIcon, PlayIcon } from "@heroicons/react/20/solid"; import { DialogClose } from "@radix-ui/react-dialog"; import { Form, useNavigation } from "@remix-run/react"; -import type { QueueItem } from "@trigger.dev/core/v3/schemas"; import { useEffect, useState } from "react"; import { cn } from "~/utils/cn"; +import type { QueueLimits } from "~/components/queues/queue-limits"; import { Button, type ButtonVariant } from "~/components/primitives/Buttons"; import { Dialog, DialogContent, DialogHeader, DialogTrigger } from "~/components/primitives/Dialog"; import { FormButtons } from "~/components/primitives/FormButtons"; +import { FormError } from "~/components/primitives/FormError"; import { Hint } from "~/components/primitives/Hint"; import { Input } from "~/components/primitives/Input"; import { InputGroup } from "~/components/primitives/InputGroup"; @@ -172,7 +173,12 @@ export function QueueOverrideConcurrencyButton({ trigger, disabled = false, }: { - queue: QueueItem & { concurrencyLimitOverridePercent: number | null }; + queue: { + id: string; + name: string; + limits: QueueLimits; + concurrencyLimitOverridePercent: number | null; + }; environmentConcurrencyLimit: number; /** How to render the dialog trigger. "menu-item" (default) is a PopoverMenuItem for row menus; * "button" is a standalone labeled button; "icon" is an icon-only button with the label in a @@ -186,14 +192,26 @@ export function QueueOverrideConcurrencyButton({ queue.concurrencyLimitOverridePercent !== null ? "percent" : "absolute" ); const [concurrencyLimit, setConcurrencyLimit] = useState( - queue.concurrencyLimit?.toString() ?? environmentConcurrencyLimit.toString() + queue.limits.perKey.current?.toString() ?? environmentConcurrencyLimit.toString() ); const [percent, setPercent] = useState( queue.concurrencyLimitOverridePercent?.toString() ?? "100" ); - const isOverridden = !!queue.concurrency?.overriddenAt; - const currentLimit = queue.concurrencyLimit ?? environmentConcurrencyLimit; + /** A row with a total bound gets the two-field dialog: each bound is overridden on its own + * and a blank field leaves that bound unchanged. Rows without a total keep the classic + * single-limit dialog (with the percent toggle). */ + const hasTotal = queue.limits.total != null; + /** Both fields start blank so blank-means-unchanged holds: a prefilled value would submit + * as an explicit override of the current value, pinning it against future code changes. + * The current values show as placeholders instead. */ + const [perKeyValue, setPerKeyValue] = useState(""); + const [totalValue, setTotalValue] = useState(""); + + const isOverridden = hasTotal + ? !!queue.limits.perKey.overriddenAt || !!queue.limits.total?.overriddenAt + : !!queue.limits.perKey.overriddenAt; + const currentLimit = queue.limits.perKey.current ?? environmentConcurrencyLimit; useEffect(() => { if (navigation.state === "loading" || navigation.state === "idle") { @@ -221,15 +239,65 @@ export function QueueOverrideConcurrencyButton({ const limitNumber = Number(concurrencyLimit); const limitOverCap = Number.isFinite(limitNumber) && limitNumber > environmentConcurrencyLimit; + const perKeyNumber = Number(perKeyValue); + const perKeyOverCap = + perKeyValue !== "" && + Number.isFinite(perKeyNumber) && + perKeyNumber > environmentConcurrencyLimit; + const perKeyInvalid = + perKeyValue !== "" && (!Number.isInteger(perKeyNumber) || perKeyNumber < 0 || perKeyOverCap); + const totalNumber = Number(totalValue); + const totalOverCap = + totalValue !== "" && Number.isFinite(totalNumber) && totalNumber > environmentConcurrencyLimit; + const totalInvalid = + totalValue !== "" && (!Number.isInteger(totalNumber) || totalNumber < 1 || totalOverCap); + + /** Cross-field check on the pair that would be in effect after submit: a blank field keeps + * its current value. A per-key limit above the total could never be reached, so creating + * that state is blocked. The API treats the bounds independently, so a pair that already + * conflicts stays editable: improvements submit with a notice instead of being trapped. */ + const resultingPerKey = perKeyValue !== "" ? perKeyNumber : queue.limits.perKey.current; + const resultingTotal = totalValue !== "" ? totalNumber : (queue.limits.total?.current ?? null); + const resultingConflict = + hasTotal && + !perKeyInvalid && + !totalInvalid && + resultingPerKey !== null && + resultingTotal !== null && + resultingPerKey > resultingTotal; + const currentConflictGap = + hasTotal && queue.limits.perKey.current !== null && queue.limits.total != null + ? queue.limits.perKey.current - queue.limits.total.current + : 0; + const resultingConflictGap = + resultingConflict && resultingPerKey !== null && resultingTotal !== null + ? resultingPerKey - resultingTotal + : 0; + const boundsConflict = + resultingConflict && !(currentConflictGap > 0 && resultingConflictGap < currentConflictGap); + const submitDisabled = disabled || - isLoading || - (mode === "percent" ? !percentValid : !concurrencyLimit || limitOverCap); + (hasTotal + ? isLoading || + perKeyInvalid || + totalInvalid || + boundsConflict || + (perKeyValue === "" && totalValue === "") + : isLoading || (mode === "percent" ? !percentValid : !concurrencyLimit || limitOverCap)); const iconLabel = isOverridden ? "Edit override" : "Override limit"; + const handleOpenChange = (open: boolean) => { + if (open) { + setPerKeyValue(""); + setTotalValue(""); + } + setIsOpen(open); + }; + return ( - + {trigger === "icon" ? ( @@ -276,7 +344,9 @@ export function QueueOverrideConcurrencyButton({ {disabled ? "You don't have permission to manage queues" - : "Give this queue its own concurrency limit instead of the environment default. Set it as a number or a percentage of the environment limit."} + : hasTotal + ? "Override this queue's per-key and total concurrency limits." + : "Give this queue its own concurrency limit instead of the environment default. Set it as a number or a percentage of the environment limit."} @@ -300,13 +370,27 @@ export function QueueOverrideConcurrencyButton({ {isOverridden ? "Edit concurrency override" : "Override concurrency limit"}
- {isOverridden ? ( + {hasTotal ? ( + isOverridden ? ( + + This queue's limits are currently overridden. Fill a field to change that limit + (blank fields stay unchanged), or remove the override to restore the limits set in + code. + + ) : ( + + Override this queue's limits. Per key caps each concurrency key's pool, and total + caps runs across all keys together. Leave a field blank to keep that limit + unchanged. + + ) + ) : isOverridden ? ( This queue's concurrency limit is currently overridden to {currentLimit}. - {typeof queue.concurrency?.base === "number" && - ` The original limit set in code was ${queue.concurrency.base}.`}{" "} + {typeof queue.limits.perKey.base === "number" && + ` The original limit set in code was ${queue.limits.perKey.base}.`}{" "} You can update the override or remove it to restore the{" "} - {typeof queue.concurrency?.base === "number" + {typeof queue.limits.perKey.base === "number" ? "limit set in code" : "environment concurrency limit"} . @@ -314,73 +398,137 @@ export function QueueOverrideConcurrencyButton({ ) : ( Override this queue's concurrency limit. The current limit is {currentLimit}, which is - set {queue.concurrencyLimit !== null ? "in code" : "by the environment"}. + set {queue.limits.perKey.current !== null ? "in code" : "by the environment"}. )}
setIsOpen(false)} className="space-y-3"> - - - -
-
- {mode === "percent" ? ( - setPercent(e.target.value)} - placeholder="100" - autoFocus - accessory={%} - /> - ) : ( - setConcurrencyLimit(e.target.value)} - placeholder={currentLimit.toString()} - autoFocus - /> - )} + + {hasTotal ? ( + <> + + + + setPerKeyValue(e.target.value)} + placeholder={queue.limits.perKey.current?.toString() ?? "No per-key limit"} + autoFocus + /> + + {perKeyOverCap + ? `Can't exceed the environment limit of ${environmentConcurrencyLimit}.` + : `The most concurrent runs each concurrency key can use at once.${ + typeof queue.limits.perKey.base === "number" + ? ` Set to ${queue.limits.perKey.base} in code.` + : "" + }`} + + + + + setTotalValue(e.target.value)} + placeholder={queue.limits.total?.current.toString()} + /> + + {totalInvalid + ? totalOverCap + ? `Can't exceed the environment limit of ${environmentConcurrencyLimit}.` + : "Enter a whole number of 1 or more." + : `The most concurrent runs across all keys together. It can't exceed the environment limit of ${environmentConcurrencyLimit}.${ + typeof queue.limits.total?.base === "number" + ? ` Set to ${queue.limits.total.base} in code.` + : "" + }`} + + + {boundsConflict ? ( + + The per-key limit ({resultingPerKey}) can't exceed the total limit ( + {resultingTotal}). + + ) : resultingConflict ? ( + + The per-key limit ({resultingPerKey}) still exceeds the total limit ( + {resultingTotal}), so only the total applies. + + ) : null} + + ) : ( + + +
+
+ {mode === "percent" ? ( + setPercent(e.target.value)} + placeholder="100" + autoFocus + accessory={%} + /> + ) : ( + setConcurrencyLimit(e.target.value)} + placeholder={currentLimit.toString()} + autoFocus + /> + )} +
+ setMode(value === "percent" ? "percent" : "absolute")} + />
- setMode(value === "percent" ? "percent" : "absolute")} - /> -
- {mode === "percent" ? ( - - {materializedFromPercent !== null - ? `${percentNumber}% = ${materializedFromPercent} concurrent ${ - materializedFromPercent === 1 ? "run" : "runs" - } of the environment's ${environmentConcurrencyLimit}. Recalculates automatically when the environment limit changes.` - : "Enter a percentage between 1 and 100."} - - ) : ( - - {limitOverCap - ? `Can't exceed the environment limit of ${environmentConcurrencyLimit}.` - : `The most concurrent runs this queue can use at once. It can't exceed the environment limit of ${environmentConcurrencyLimit}.`} - - )} - + {mode === "percent" ? ( + + {materializedFromPercent !== null + ? `${percentNumber}% = ${materializedFromPercent} concurrent ${ + materializedFromPercent === 1 ? "run" : "runs" + } of the environment's ${environmentConcurrencyLimit}. Recalculates automatically when the environment limit changes.` + : "Enter a percentage between 1 and 100."} + + ) : ( + + {limitOverCap + ? `Can't exceed the environment limit of ${environmentConcurrencyLimit}.` + : `The most concurrent runs this queue can use at once. It can't exceed the environment limit of ${environmentConcurrencyLimit}.`} + + )} + + )} + + } + content={ + displayName.startsWith(TASK_PREFIX) + ? `This is the inline concurrency limit of your "${displayName.slice( + TASK_PREFIX.length + )}" task` + : "This is a named concurrency limit declared in your code." + } + /> + {displayName} + + ); + } + return ( {type === "task" ? ( diff --git a/apps/webapp/app/env.server.ts b/apps/webapp/app/env.server.ts index 63b966e8b05..2dc04ded4c7 100644 --- a/apps/webapp/app/env.server.ts +++ b/apps/webapp/app/env.server.ts @@ -1440,6 +1440,8 @@ const EnvironmentSchema = z RUN_ENGINE_RUN_QUEUE_LOG_LEVEL: z .enum(["log", "error", "warn", "info", "debug"]) .default("info"), + RUN_ENGINE_TOTAL_CONCURRENCY_LIMITS_ENABLED: z.string().default("1"), + RUN_ENGINE_QUEUE_GATES_ENABLED: z.string().default("0"), RUN_ENGINE_TREAT_PRODUCTION_EXECUTION_STALLS_AS_OOM: z.string().default("0"), RUN_ENGINE_READ_REPLICA_SNAPSHOTS_SINCE_ENABLED: z.string().default("0"), RUN_ENGINE_SNAPSHOTS_SINCE_REPLICA_RETRY_MIN_MS: z.coerce.number().int().default(50), diff --git a/apps/webapp/app/models/queueMutation.server.ts b/apps/webapp/app/models/queueMutation.server.ts index 5e4dbcd90e4..5cc972a050b 100644 --- a/apps/webapp/app/models/queueMutation.server.ts +++ b/apps/webapp/app/models/queueMutation.server.ts @@ -62,12 +62,109 @@ export async function handleQueueMutationAction({ } case "queue-override": { const friendlyId = formData.get("friendlyId"); - const mode = formData.get("mode") === "percent" ? "percent" : "absolute"; + const mode = + formData.get("mode") === "percent" + ? "percent" + : formData.get("mode") === "bounds" + ? "bounds" + : "absolute"; if (!friendlyId) { return redirectWithErrorMessage(redirectPath, request, "Queue ID is required"); } + if (mode === "bounds") { + const user = await getUserById(userId); + if (!user) { + return redirectWithErrorMessage(redirectPath, request, "User not found"); + } + + const perKeyRaw = formData.get("perKeyLimit")?.toString().trim() || null; + const totalRaw = formData.get("totalLimit")?.toString().trim() || null; + + if (perKeyRaw === null && totalRaw === null) { + return redirectWithErrorMessage( + redirectPath, + request, + "Enter a per-key limit, a total limit, or both" + ); + } + + const perKey = perKeyRaw === null ? null : Number(perKeyRaw); + if ( + perKey !== null && + (!Number.isInteger(perKey) || perKey < 0 || perKey > environment.maximumConcurrencyLimit) + ) { + return redirectWithErrorMessage( + redirectPath, + request, + `Per-key limit must be a whole number between 0 and the environment limit of ${environment.maximumConcurrencyLimit}` + ); + } + + const total = totalRaw === null ? null : Number(totalRaw); + if ( + total !== null && + (!Number.isInteger(total) || total < 1 || total > environment.maximumConcurrencyLimit) + ) { + return redirectWithErrorMessage( + redirectPath, + request, + `Total limit must be a whole number between 1 and the environment limit of ${environment.maximumConcurrencyLimit}` + ); + } + + if (perKey !== null) { + const result = await concurrencySystem.queues.overrideQueueConcurrencyLimit( + environment, + friendlyId.toString(), + { limit: perKey }, + user + ); + if (!result.isOk()) { + const error = result.error; + const message = + "message" in error && typeof error.message === "string" + ? error.message + : "Failed to override the per-key limit"; + return redirectWithErrorMessage(redirectPath, request, message); + } + } + + if (total !== null) { + const result = await concurrencySystem.queues.overrideTotalConcurrencyLimit( + environment, + friendlyId.toString(), + total, + user + ); + if (!result.isOk()) { + const error = result.error; + const message = + "message" in error && typeof error.message === "string" + ? error.message + : "Failed to override the total limit"; + return redirectWithErrorMessage( + redirectPath, + request, + perKey !== null + ? `The per-key limit was overridden, but the total limit failed: ${message}` + : message + ); + } + } + + return redirectWithSuccessMessage( + redirectPath, + request, + perKey !== null && total !== null + ? "Per-key and total limits overridden" + : perKey !== null + ? "Per-key limit overridden" + : "Total limit overridden" + ); + } + // The dialog submits either a `percent` of the environment limit or an absolute `limit`, // depending on the unit toggle. Build the matching override shape for the service. let override: number | { limit: number } | { percent: number }; @@ -137,12 +234,14 @@ export async function handleQueueMutationAction({ return redirectWithErrorMessage(redirectPath, request, "Queue ID is required"); } + const isBoundsScope = formData.get("scope") === "bounds"; + const result = await concurrencySystem.queues.resetConcurrencyLimit( environment, friendlyId.toString() ); - if (!result.isOk()) { + if (!result.isOk() && !(isBoundsScope && result.error.type === "queue_not_overridden")) { return redirectWithErrorMessage( redirectPath, request, @@ -150,6 +249,22 @@ export async function handleQueueMutationAction({ ); } + if (isBoundsScope) { + const totalResult = await concurrencySystem.queues.resetTotalConcurrencyLimit( + environment, + friendlyId.toString() + ); + if (!totalResult.isOk() && totalResult.error.type !== "queue_not_overridden") { + return redirectWithErrorMessage( + redirectPath, + request, + result.isOk() + ? "The per-key limit was reset, but resetting the total limit failed" + : "Failed to reset the total limit" + ); + } + } + return redirectWithSuccessMessage(redirectPath, request, "Queue concurrency limit reset"); } default: diff --git a/apps/webapp/app/presenters/v3/BuiltInDashboards.server.ts b/apps/webapp/app/presenters/v3/BuiltInDashboards.server.ts index d831568248d..7c609322ab2 100644 --- a/apps/webapp/app/presenters/v3/BuiltInDashboards.server.ts +++ b/apps/webapp/app/presenters/v3/BuiltInDashboards.server.ts @@ -634,7 +634,7 @@ const queuesDashboard: BuiltInDashboard = { "t-pressure": { title: "Queue pressure", query: "", display: { type: "title" } }, pressure: { title: "Queue pressure", - query: `SELECT queue,\n argMax(max_running, bucket_start) AS running,\n argMax(max_queued, bucket_start) AS queued,\n argMax(max_limit, bucket_start) AS limit,\n running + queued AS demand,\n max(max_queued) AS peak_queued,\n sum(throttled_count) AS throttled,\n multiIf(running >= limit AND queued > 0, 'queue-limited', queued > 0, 'backlogged', 'healthy') AS status\nFROM queue_metrics\nGROUP BY queue\nORDER BY peak_queued DESC`, + query: `SELECT queue,\n argMax(max_running, bucket_start) AS running,\n argMax(max_queued, bucket_start) AS queued,\n argMax(max_limit, bucket_start) AS limit,\n running + queued AS demand,\n max(max_queued) AS peak_queued,\n sum(throttled_count) AS throttled,\n multiIf(running >= limit AND queued > 0, 'queue-limited', queued > 0, 'backlogged', 'healthy') AS status\nFROM concurrency_metrics\nGROUP BY queue\nORDER BY peak_queued DESC`, display: { type: "table", prettyFormatting: true, @@ -644,7 +644,7 @@ const queuesDashboard: BuiltInDashboard = { "t-trends": { title: "Per-queue trends", query: "", display: { type: "title" } }, "running-q": { title: "Running by queue", - query: `SELECT timeBucket() AS t, queue, max(max_running) AS running\nFROM queue_metrics\nGROUP BY t, queue\nORDER BY t`, + query: `SELECT timeBucket() AS t, queue, max(max_running) AS running\nFROM concurrency_metrics\nGROUP BY t, queue\nORDER BY t`, // Grouped gauge: carry each queue's running across idle buckets (per-group LOCF). fillGaps: true, display: { @@ -661,7 +661,7 @@ const queuesDashboard: BuiltInDashboard = { }, "queued-q": { title: "Queue depth (backlog) by queue", - query: `SELECT timeBucket() AS t, queue, max(max_queued) AS queued\nFROM queue_metrics\nGROUP BY t, queue\nORDER BY t`, + query: `SELECT timeBucket() AS t, queue, max(max_queued) AS queued\nFROM concurrency_metrics\nGROUP BY t, queue\nORDER BY t`, // Grouped gauge: carry each queue's backlog across idle buckets (per-group LOCF). fillGaps: true, display: { @@ -678,7 +678,7 @@ const queuesDashboard: BuiltInDashboard = { }, "throttled-q": { title: "Throttled buckets by queue", - query: `SELECT timeBucket() AS t, queue, sum(throttled_count) AS throttled\nFROM queue_metrics\nGROUP BY t, queue\nORDER BY t`, + query: `SELECT timeBucket() AS t, queue, sum(throttled_count) AS throttled\nFROM concurrency_metrics\nGROUP BY t, queue\nORDER BY t`, // Grouped counter: per-group zero-fill so idle buckets read 0, not a gap. fillGaps: true, display: { @@ -697,7 +697,7 @@ const queuesDashboard: BuiltInDashboard = { title: "Enqueued vs started", // Counter states merge per queue, then sum outside: a single merge across queues // mixes unrelated odometers and returns wrong totals. - query: `SELECT t, sum(enq) AS enqueued, sum(st) AS started\nFROM (\n SELECT timeBucket() AS t, queue,\n deltaSumTimestampMerge(enqueue_delta) AS enq,\n deltaSumTimestampMerge(started_delta) AS st\n FROM queue_metrics\n GROUP BY t, queue\n)\nGROUP BY t\nORDER BY t`, + query: `SELECT t, sum(enq) AS enqueued, sum(st) AS started\nFROM (\n SELECT timeBucket() AS t, queue,\n deltaSumTimestampMerge(enqueue_delta) AS enq,\n deltaSumTimestampMerge(started_delta) AS st\n FROM concurrency_metrics\n GROUP BY t, queue\n)\nGROUP BY t\nORDER BY t`, display: { type: "chart", chartType: "line", diff --git a/apps/webapp/app/presenters/v3/QueueListPresenter.server.ts b/apps/webapp/app/presenters/v3/QueueListPresenter.server.ts index 0dc3daa9856..05353f2ccdb 100644 --- a/apps/webapp/app/presenters/v3/QueueListPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/QueueListPresenter.server.ts @@ -7,9 +7,16 @@ import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactoryInstan import { logger } from "~/services/logger.server"; import { engine } from "~/v3/runEngine.server"; import { BasePresenter } from "./basePresenter.server"; -import { toQueueItem } from "./QueueRetrievePresenter.server"; +import { toQueueItem, toQueueLimits } from "./QueueRetrievePresenter.server"; +import type { QueueLimits } from "~/components/queues/queue-limits"; -type QueueListEngine = Pick; +type QueueListEngine = Pick< + RunEngine, + | "lengthOfQueues" + | "currentConcurrencyOfQueues" + | "totalConcurrencyOfQueues" + | "gateQueuedCountOfQueues" +>; export const QUEUE_LIST_DEFAULT_ITEMS_PER_PAGE = 25; const MAX_ITEMS_PER_PAGE = 100; @@ -34,8 +41,13 @@ const queueListSelect = { concurrencyLimitOverriddenAt: true, concurrencyLimitOverriddenBy: true, concurrencyLimitOverridePercent: true, + totalConcurrencyLimit: true, + totalConcurrencyLimitBase: true, + totalConcurrencyLimitOverriddenAt: true, type: true, paused: true, + role: true, + concurrencyVersion: true, } satisfies Prisma.TaskQueueSelect; type QueueListRow = Prisma.TaskQueueGetPayload<{ select: typeof queueListSelect }>; @@ -44,6 +56,12 @@ type QueueListRow = Prisma.TaskQueueGetPayload<{ select: typeof queueListSelect // schema (that's a public contract), so we surface it as an extra field on the list item. type QueueListItem = ReturnType & { concurrencyLimitOverridePercent: number | null; + /** "queue" rows wait and order runs; "limit" rows are named concurrency limits. */ + kind: "queue" | "limit"; + /** V2 rows hold the new perKey/total vocabulary in their limit columns. */ + concurrencyVersion: "V1" | "V2"; + /** The row's configured bounds, dashboard-only (the public shape hides them on V2). */ + limits: QueueLimits; }; type QueueListPagination = @@ -64,21 +82,49 @@ function formatClickhouseDateTime(date: Date): string { function buildQueueListWhere( environmentId: string, query: string | undefined, - type: "task" | "custom" | undefined + type: "task" | "custom" | undefined, + includeLimits: boolean ): Prisma.TaskQueueWhereInput { const trimmedQuery = query?.trim(); - return { + const common = { runtimeEnvironmentId: environmentId, - version: "V2", + version: "V2" as const, name: trimmedQuery ? { contains: trimmedQuery, - mode: "insensitive", + mode: "insensitive" as const, } : undefined, type: type ? typeToDBQueueType[type] : undefined, }; + + /** Only the dashboard interleaves named limits, and the type filter names queue + * shapes, so either condition scopes the list to queue rows; the public queues + * API always stays queue-only. Boundless rows in the anonymous limit/task/ + * namespace are retired (their inline limit moved onto the task's own queue) + * and stay hidden; boundless NAMED limits are real uncapped rows and show. */ + if (includeLimits && !type) { + return { + ...common, + OR: [ + { role: "QUEUE" as const }, + { + role: "LIMIT" as const, + OR: [ + { name: { not: { startsWith: "limit/task/" } } }, + { concurrencyLimit: { not: null } }, + { totalConcurrencyLimit: { not: null } }, + ], + }, + ], + }; + } + + return { + ...common, + role: "QUEUE" as const, + }; } export class QueueListPresenter extends BasePresenter { @@ -102,6 +148,7 @@ export class QueueListPresenter extends BasePresenter { page, type, sort = "name", + includeLimits = false, }: { environment: AuthenticatedEnvironment; query?: string; @@ -109,13 +156,21 @@ export class QueueListPresenter extends BasePresenter { perPage?: number; type?: "task" | "custom"; sort?: QueueListSort; + includeLimits?: boolean; }): Promise { const hasFilters = Boolean(query?.trim()) || type !== undefined; if (sort !== "name") { // Ranking is additive: any failure or unsupported input falls back to name order. try { - const ranked = await this.getRankedQueues(environment, query, page, type, sort); + const ranked = await this.getRankedQueues( + environment, + query, + page, + type, + sort, + includeLimits + ); if (ranked) { return ranked; } @@ -125,7 +180,13 @@ export class QueueListPresenter extends BasePresenter { } if (hasFilters) { - const { queues, hasMore } = await this.getFilteredQueues(environment, query, page, type); + const { queues, hasMore } = await this.getFilteredQueues( + environment, + query, + page, + type, + includeLimits + ); return { queues, @@ -139,11 +200,11 @@ export class QueueListPresenter extends BasePresenter { } const totalQueues = await this._replica.taskQueue.count({ - where: buildQueueListWhere(environment.id, query, type), + where: buildQueueListWhere(environment.id, query, type, includeLimits), }); return { - queues: await this.getUnfilteredQueues(environment, page, type), + queues: await this.getUnfilteredQueues(environment, page, type, includeLimits), pagination: { mode: "unfiltered" as const, currentPage: page, @@ -164,7 +225,8 @@ export class QueueListPresenter extends BasePresenter { query: string | undefined, page: number, type: "task" | "custom" | undefined, - sort: Exclude + sort: Exclude, + includeLimits: boolean ) { if (type !== undefined) { return null; @@ -213,7 +275,7 @@ export class QueueListPresenter extends BasePresenter { return null; } - const where = buildQueueListWhere(environment.id, query, type); + const where = buildQueueListWhere(environment.id, query, type, includeLimits); const totalQueues = await this._replica.taskQueue.count({ where }); let rankedPageQueues: QueueListRow[] = []; @@ -284,10 +346,11 @@ export class QueueListPresenter extends BasePresenter { environment: AuthenticatedEnvironment, query: string | undefined, page: number, - type: "task" | "custom" | undefined + type: "task" | "custom" | undefined, + includeLimits: boolean ) { const queues = await this._replica.taskQueue.findMany({ - where: buildQueueListWhere(environment.id, query, type), + where: buildQueueListWhere(environment.id, query, type, includeLimits), select: queueListSelect, orderBy: { orderableName: "asc", @@ -307,10 +370,11 @@ export class QueueListPresenter extends BasePresenter { private async getUnfilteredQueues( environment: AuthenticatedEnvironment, page: number, - type: "task" | "custom" | undefined + type: "task" | "custom" | undefined, + includeLimits: boolean ) { const queues = await this._replica.taskQueue.findMany({ - where: buildQueueListWhere(environment.id, undefined, type), + where: buildQueueListWhere(environment.id, undefined, type, includeLimits), select: queueListSelect, orderBy: { orderableName: "asc", @@ -333,20 +397,49 @@ export class QueueListPresenter extends BasePresenter { concurrencyLimitOverriddenAt: Date | null; concurrencyLimitOverriddenBy: string | null; concurrencyLimitOverridePercent: Prisma.Decimal | null; + totalConcurrencyLimit: number | null; + totalConcurrencyLimitBase: number | null; + totalConcurrencyLimitOverriddenAt: Date | null; type: TaskQueueType; paused: boolean; + role: "QUEUE" | "LIMIT"; + concurrencyVersion: "V1" | "V2"; }[] ): Promise { - const [queuedByQueue, runningByQueue] = await Promise.all([ - this.engineClient.lengthOfQueues( - environment, - queues.map((q) => q.name) - ), - this.engineClient.currentConcurrencyOfQueues( - environment, - queues.map((q) => q.name) - ), - ]); + const queueRows = queues.filter((q) => q.role === "QUEUE"); + const limitRows = queues.filter((q) => q.role === "LIMIT"); + /** + * Queue rows read their zset length and home concurrency; limit rows read the + * group set (every holder, keyed or keyless) as running and the per-gate queued + * counter as queued. The group read also serves queue rows with a total cap. + */ + const rowsWithGroupRead = [ + ...queueRows.filter((q) => q.totalConcurrencyLimit !== null), + ...limitRows, + ]; + const [queuedByQueue, runningByQueue, totalRunningByQueue, gateQueuedByQueue] = + await Promise.all([ + this.engineClient.lengthOfQueues( + environment, + queueRows.map((q) => q.name) + ), + this.engineClient.currentConcurrencyOfQueues( + environment, + queueRows.map((q) => q.name) + ), + rowsWithGroupRead.length > 0 + ? this.engineClient.totalConcurrencyOfQueues( + environment, + rowsWithGroupRead.map((q) => q.name) + ) + : Promise.resolve({} as Record), + limitRows.length > 0 + ? this.engineClient.gateQueuedCountOfQueues( + environment, + limitRows.map((q) => q.name) + ) + : Promise.resolve({} as Record), + ]); // Manually "join" the overridden users because there is no way to implement the relationship // in prisma without adding a foreign key constraint @@ -359,26 +452,45 @@ export class QueueListPresenter extends BasePresenter { const overriddenByMap = new Map(overriddenByUsers.map((u) => [u.id, u])); - return queues.map((queue) => ({ - ...toQueueItem({ - friendlyId: queue.friendlyId, - name: queue.name, - type: queue.type, - running: runningByQueue[queue.name] ?? 0, - queued: queuedByQueue[queue.name] ?? 0, - concurrencyLimit: queue.concurrencyLimit ?? null, - concurrencyLimitBase: queue.concurrencyLimitBase ?? null, - concurrencyLimitOverriddenAt: queue.concurrencyLimitOverriddenAt ?? null, - concurrencyLimitOverriddenBy: queue.concurrencyLimitOverriddenBy - ? (overriddenByMap.get(queue.concurrencyLimitOverriddenBy) ?? null) - : null, - paused: queue.paused, - }), - // Prisma returns Decimal; the client only needs a plain number (null for absolute overrides). - concurrencyLimitOverridePercent: - queue.concurrencyLimitOverridePercent !== null - ? Number(queue.concurrencyLimitOverridePercent) - : null, - })); + return queues.map((queue) => { + const overriddenByUser = queue.concurrencyLimitOverriddenBy + ? (overriddenByMap.get(queue.concurrencyLimitOverriddenBy) ?? null) + : null; + return { + ...toQueueItem({ + friendlyId: queue.friendlyId, + name: queue.name, + type: queue.type, + version: queue.concurrencyVersion, + running: + queue.role === "LIMIT" + ? (totalRunningByQueue[queue.name] ?? 0) + : (runningByQueue[queue.name] ?? 0), + queued: + queue.role === "LIMIT" + ? (gateQueuedByQueue[queue.name] ?? 0) + : (queuedByQueue[queue.name] ?? 0), + concurrencyLimit: queue.concurrencyLimit ?? null, + concurrencyLimitBase: queue.concurrencyLimitBase ?? null, + concurrencyLimitOverriddenAt: queue.concurrencyLimitOverriddenAt ?? null, + concurrencyLimitOverriddenBy: overriddenByUser, + paused: queue.paused, + }), + // Prisma returns Decimal; the client only needs a plain number (null for absolute overrides). + concurrencyLimitOverridePercent: + queue.concurrencyLimitOverridePercent !== null + ? Number(queue.concurrencyLimitOverridePercent) + : null, + kind: queue.role === "LIMIT" ? ("limit" as const) : ("queue" as const), + concurrencyVersion: queue.concurrencyVersion, + limits: toQueueLimits(queue, { + totalRunning: + queue.totalConcurrencyLimit !== null ? (totalRunningByQueue[queue.name] ?? 0) : null, + overriddenByName: overriddenByUser + ? (overriddenByUser.displayName ?? overriddenByUser.name ?? null) + : null, + }), + }; + }); } } diff --git a/apps/webapp/app/presenters/v3/QueueRetrievePresenter.server.ts b/apps/webapp/app/presenters/v3/QueueRetrievePresenter.server.ts index f6918394e5c..b05d3ed5082 100644 --- a/apps/webapp/app/presenters/v3/QueueRetrievePresenter.server.ts +++ b/apps/webapp/app/presenters/v3/QueueRetrievePresenter.server.ts @@ -1,5 +1,11 @@ import { assertExhaustive } from "@trigger.dev/core"; -import { type Prettify, type QueueItem, type RetrieveQueueParam } from "@trigger.dev/core/v3"; +import { + QueueItem as QueueItemSchema, + type Prettify, + type QueueItem, + type RetrieveQueueParam, +} from "@trigger.dev/core/v3"; +import type { QueueLimits } from "~/components/queues/queue-limits"; import { type PrismaClientOrTransaction, type TaskQueue, @@ -24,6 +30,8 @@ export async function getQueue( environment: AuthenticatedEnvironment, queue: RetrieveQueueParam ) { + const role = "QUEUE" as const; + if (typeof queue === "string") { return joinQueueWithUser( prismaClient, @@ -31,6 +39,7 @@ export async function getQueue( where: { friendlyId: queue, runtimeEnvironmentId: environment.id, + role, }, }) ); @@ -44,6 +53,7 @@ export async function getQueue( where: { name: queueName, runtimeEnvironmentId: environment.id, + role, }, }) ); @@ -90,9 +100,15 @@ export class QueueRetrievePresenter extends BasePresenter { const results = await Promise.all([ engine.lengthOfQueues(environment, [queue.name]), engine.currentConcurrencyOfQueues(environment, [queue.name]), + queue.totalConcurrencyLimit != null + ? engine.totalConcurrencyOfQueues(environment, [queue.name]) + : undefined, ]); - // Transform queues to include running and queued counts + /** The returned queue = the public QueueItem fields plus dashboard extras + * (percent override source, configured bounds); the public API routes strip + * the extras via `toPublicQueueItem`. Prisma returns Decimal for the + * percent; the client only needs a plain number (null for absolute). */ return { success: true as const, queue: { @@ -100,6 +116,7 @@ export class QueueRetrievePresenter extends BasePresenter { friendlyId: queue.friendlyId, name: queue.name, type: queue.type, + version: queue.concurrencyVersion, running: results[1]?.[queue.name] ?? 0, queued: results[0]?.[queue.name] ?? 0, concurrencyLimit: queue.concurrencyLimit ?? null, @@ -108,19 +125,71 @@ export class QueueRetrievePresenter extends BasePresenter { concurrencyLimitOverriddenBy: queue.concurrencyLimitOverriddenBy ?? null, paused: queue.paused, }), - // The percent source-of-truth for percent-based overrides isn't part of the shared - // `QueueItem` schema (that's a public contract), so we surface it as an extra field on - // the returned queue — mirroring QueueListPresenter. Prisma returns Decimal; the client - // only needs a plain number (null for absolute overrides). concurrencyLimitOverridePercent: queue.concurrencyLimitOverridePercent !== null ? Number(queue.concurrencyLimitOverridePercent) : null, + limits: toQueueLimits(queue, { + totalRunning: + queue.totalConcurrencyLimit != null ? (results[2]?.[queue.name] ?? 0) : null, + overriddenByName: toQueueConcurrencyOverriddenBy( + queue.concurrencyLimitOverriddenBy ?? null + ), + }), }, }; } } +/** + * The dashboard-only configured bounds of a row, independent of the public + * shape's version discrimination (V2 queues expose no queue-level concurrency + * publicly, but the dashboard still shows what's configured). + */ +export function toQueueLimits( + row: { + concurrencyLimit: number | null; + concurrencyLimitBase: number | null; + concurrencyLimitOverriddenAt: Date | null; + totalConcurrencyLimit: number | null; + totalConcurrencyLimitBase: number | null; + totalConcurrencyLimitOverriddenAt: Date | null; + }, + extras: { totalRunning: number | null; overriddenByName: string | null } +): QueueLimits { + return { + perKey: { + current: row.concurrencyLimit, + base: row.concurrencyLimitOverriddenAt ? row.concurrencyLimitBase : row.concurrencyLimit, + override: row.concurrencyLimitOverriddenAt ? row.concurrencyLimit : null, + overriddenAt: row.concurrencyLimitOverriddenAt, + overriddenBy: extras.overriddenByName, + }, + total: + row.totalConcurrencyLimit !== null + ? { + current: row.totalConcurrencyLimit, + base: row.totalConcurrencyLimitOverriddenAt + ? row.totalConcurrencyLimitBase + : row.totalConcurrencyLimit, + override: row.totalConcurrencyLimitOverriddenAt ? row.totalConcurrencyLimit : null, + overriddenAt: row.totalConcurrencyLimitOverriddenAt, + running: extras.totalRunning, + } + : null, + }; +} + +/** + * The public API wire shape: the `QueueItem` contract exactly (dashboard extras + * stripped by the schema parse) plus the legacy field older clients require. + */ +export function toPublicQueueItem(item: QueueItem): QueueItem & { + releaseConcurrencyOnWaitpoint: boolean; +} { + return { ...QueueItemSchema.parse(item), releaseConcurrencyOnWaitpoint: true }; +} + export function queueTypeFromType(type: TaskQueueType) { switch (type) { case "NAMED": @@ -133,14 +202,17 @@ export function queueTypeFromType(type: TaskQueueType) { } /** - * Converts raw queue data into a standardized QueueItem format - * @param data Raw queue data containing required queue properties - * @returns A validated QueueItem object + * Converts raw queue data into the public QueueItem shape. The queue's version + * discriminates it: V1 carries the queue's own limit and override state; V2 + * carries neither (a V2 queue is only the line — concurrency lives on the task + * `concurrency` option and the `concurrencyLimits` surface), with + * `concurrencyLimit` kept as null so older clients keep parsing. */ export function toQueueItem(data: { friendlyId: string; name: string; type: TaskQueueType; + version: "V1" | "V2"; running: number; queued: number; concurrencyLimit: number | null; @@ -149,7 +221,7 @@ export function toQueueItem(data: { concurrencyLimitOverriddenBy: User | null; paused: boolean; }): QueueItem & { releaseConcurrencyOnWaitpoint: boolean } { - return { + const common = { id: data.friendlyId, //remove the task/ prefix if it exists name: data.name.replace(/^task\//, ""), @@ -157,6 +229,21 @@ export function toQueueItem(data: { running: data.running, queued: data.queued, paused: data.paused, + // TODO: This needs to be removed but keeping this here for now to avoid breaking existing clients + releaseConcurrencyOnWaitpoint: true, + }; + + if (data.version === "V2") { + return { + ...common, + version: "V2" as const, + concurrencyLimit: null, + }; + } + + return { + ...common, + version: "V1" as const, concurrencyLimit: data.concurrencyLimit, concurrency: { current: data.concurrencyLimit, @@ -165,8 +252,6 @@ export function toQueueItem(data: { overriddenBy: toQueueConcurrencyOverriddenBy(data.concurrencyLimitOverriddenBy), overriddenAt: data.concurrencyLimitOverriddenAt, }, - // TODO: This needs to be removed but keeping this here for now to avoid breaking existing clients - releaseConcurrencyOnWaitpoint: true, }; } diff --git a/apps/webapp/app/presenters/v3/reports/health/flow.ts b/apps/webapp/app/presenters/v3/reports/health/flow.ts index ce6ba724527..e1d2e77654f 100644 --- a/apps/webapp/app/presenters/v3/reports/health/flow.ts +++ b/apps/webapp/app/presenters/v3/reports/health/flow.ts @@ -116,7 +116,7 @@ export function interpretFlow(metrics: Metric[], input: HealthInput): Finding { exclusions: [], observations: finishedPerMin > 0 ? [{ code: "not_workers_platform", evidence: { finishedPerMin } }] : [], - recommendation: { code: "raise_env_limit", link: "concurrency" }, + recommendation: { code: "raise_env_limit", link: "concurrency-limits" }, usesAttribution: true, }; } else if (ev.throttledShare >= t.throttledShare && !pinned) { diff --git a/apps/webapp/app/presenters/v3/reports/health/health-data.ts b/apps/webapp/app/presenters/v3/reports/health/health-data.ts index 3c14b498e57..de2cc01b4cd 100644 --- a/apps/webapp/app/presenters/v3/reports/health/health-data.ts +++ b/apps/webapp/app/presenters/v3/reports/health/health-data.ts @@ -212,7 +212,7 @@ function queueWorstQuery(): string { return `SELECT queue AS name, argMax(max_queued, bucket_start) AS latest_queued -FROM queue_metrics +FROM concurrency_metrics GROUP BY queue ORDER BY latest_queued DESC LIMIT 20`; @@ -228,7 +228,7 @@ FROM ( SELECT deltaSumTimestampMerge(dlq_delta) AS dlq, argMax(max_queued, bucket_start) AS latest_queued - FROM queue_metrics + FROM concurrency_metrics GROUP BY queue )`; } diff --git a/apps/webapp/app/presenters/v3/reports/report-registry.ts b/apps/webapp/app/presenters/v3/reports/report-registry.ts index 23cc178befb..d8b3bc4e798 100644 --- a/apps/webapp/app/presenters/v3/reports/report-registry.ts +++ b/apps/webapp/app/presenters/v3/reports/report-registry.ts @@ -4,7 +4,7 @@ import { loadHealthInput } from "./health/health-data"; import { type ReportViewModel } from "./report-view-model"; /** A query table a report may read. Same table names the query API authorizes against. */ -export type ReportQueryTable = "runs" | "env_metrics" | "queue_metrics"; +export type ReportQueryTable = "runs" | "env_metrics" | "concurrency_metrics"; export type ReportLoader = { /** Authorization metadata: the route derives its per-table JWT scope check from this. */ @@ -19,7 +19,7 @@ function defineReport(loader: ReportLoader): ReportLoader> = { health: defineReport({ - tables: ["runs", "env_metrics", "queue_metrics"], + tables: ["runs", "env_metrics", "concurrency_metrics"], load: (env, period) => loadHealthInput(env, period), interpret: interpretHealth, }), diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.concurrency-limits/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.concurrency-limits/route.tsx new file mode 100644 index 00000000000..5faab19d5c1 --- /dev/null +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.concurrency-limits/route.tsx @@ -0,0 +1,923 @@ +import { getFormProps, getInputProps, useForm } from "@conform-to/react"; +import { parseWithZod } from "@conform-to/zod/v4"; +import { + ArrowDownIcon, + EnvelopeIcon, + ExclamationTriangleIcon, + InformationCircleIcon, +} from "@heroicons/react/20/solid"; +import { DialogClose } from "@radix-ui/react-dialog"; +import { Form, useActionData, useNavigation, useSearchParams } from "@remix-run/react"; +import { json, type ActionFunctionArgs, type LoaderFunctionArgs } from "@remix-run/server-runtime"; +import { tryCatch } from "@trigger.dev/core"; +import { useEffect, useState } from "react"; +import { typedjson, useTypedLoaderData } from "remix-typedjson"; +import simplur from "simplur"; +import { z } from "zod"; +import { AdminDebugTooltip } from "~/components/admin/debugTooltip"; +import { CopyableText } from "~/components/primitives/CopyableText"; +import { EnvironmentCombo } from "~/components/environments/EnvironmentLabel"; +import { Feedback } from "~/components/Feedback"; +import { + MainHorizontallyCenteredContainer, + PageBody, + PageContainer, +} from "~/components/layout/AppLayout"; +import { Button, LinkButton } from "~/components/primitives/Buttons"; +import { Dialog, DialogContent, DialogHeader, DialogTrigger } from "~/components/primitives/Dialog"; +import { Fieldset } from "~/components/primitives/Fieldset"; +import { FormButtons } from "~/components/primitives/FormButtons"; +import { FormError } from "~/components/primitives/FormError"; +import { Header2, Header3 } from "~/components/primitives/Headers"; +import { Input } from "~/components/primitives/Input"; +import { InputGroup } from "~/components/primitives/InputGroup"; +import { InputNumberStepper } from "~/components/primitives/InputNumberStepper"; +import { Label } from "~/components/primitives/Label"; +import { NavBar, PageAccessories, PageTitle } from "~/components/primitives/PageHeader"; +import { Paragraph } from "~/components/primitives/Paragraph"; +import * as Property from "~/components/primitives/PropertyTable"; +import { SpinnerWhite } from "~/components/primitives/Spinner"; +import { + Table, + TableBody, + TableCell, + TableHeader, + TableHeaderCell, + TableRow, +} from "~/components/primitives/Table"; +import { InfoIconTooltip } from "~/components/primitives/Tooltip"; +import { useFeatures } from "~/hooks/useFeatures"; +import { useOrganization } from "~/hooks/useOrganizations"; +import { useShowSelfServe } from "~/hooks/useShowSelfServe"; +import { redirectWithErrorMessage, redirectWithSuccessMessage } from "~/models/message.server"; +import { findProjectBySlug } from "~/models/project.server"; +import { + ManageConcurrencyPresenter, + type ConcurrencyResult, + type EnvironmentWithConcurrency, +} from "~/presenters/v3/ManageConcurrencyPresenter.server"; +import { + getCurrentPlan, + getPlans, + getSelfServePurchaseBlockReason, +} from "~/services/platform.v3.server"; +import { textLinkClassName } from "~/components/primitives/TextLink"; +import { rbac } from "~/services/rbac.server"; +import { requireUserId } from "~/services/session.server"; +import { cn } from "~/utils/cn"; +import { formatCurrency, formatNumber } from "~/utils/numberFormatter"; +import { isPaidAddOnPurchase } from "~/utils/paidAddOnPermissions"; +import { concurrencyLimitsPath, EnvironmentParamSchema, v3BillingPath } from "~/utils/pathBuilder"; +import { AllocateConcurrencyService } from "~/v3/services/allocateConcurrency.server"; +import { SetConcurrencyAddOnService } from "~/v3/services/setConcurrencyAddOn.server"; +import { useCurrentPlan } from "../_app.orgs.$organizationSlug/route"; +import { sectionAgentPageContext } from "~/components/dashboard-agent/suggested-prompts"; +import type { Handle } from "~/utils/handle"; + +export const handle: Handle = { + agentPageContext: () => sectionAgentPageContext("concurrency"), +}; +import { pageMeta } from "~/utils/pageTitle"; + +export const meta = pageMeta("Manage concurrency"); + +export const loader = async ({ request, params }: LoaderFunctionArgs) => { + const userId = await requireUserId(request); + const { + organizationSlug, + projectParam, + envParam: _envParam, + } = EnvironmentParamSchema.parse(params); + + const project = await findProjectBySlug(organizationSlug, projectParam, userId); + if (!project) { + throw new Response(undefined, { + status: 404, + statusText: "Project not found", + }); + } + + const auth = await rbac.authenticateSession(request, { + userId, + organizationId: project.organizationId, + }); + const canManageBilling = auth.ok && auth.ability.can("manage", { type: "billing" }); + + const presenter = new ManageConcurrencyPresenter(); + const [error, result] = await tryCatch( + presenter.call({ + userId: userId, + projectId: project.id, + organizationId: project.organizationId, + }) + ); + + if (error) { + throw new Response(undefined, { + status: 400, + statusText: error.message, + }); + } + + const plans = await tryCatch(getPlans()); + if (!plans) { + throw new Response(null, { status: 404, statusText: "Plans not found" }); + } + + return typedjson({ ...result, canManageBilling }); +}; + +const FormSchema = z.discriminatedUnion("action", [ + z.object({ + action: z.enum(["purchase"]), + amount: z.coerce.number().min(0, "Amount must be 0 or more"), + }), + z.object({ + action: z.enum(["quota-increase"]), + amount: z.coerce.number().min(1, "Amount must be greater than 0"), + }), + z.object({ + action: z.enum(["allocate"]), + // It will only update environments that are passed in + environments: z.array( + z.object({ + id: z.string(), + amount: z.coerce.number().min(0, "Amount must be 0 or more"), + }) + ), + }), +]); + +export const action = async ({ request, params }: ActionFunctionArgs) => { + const userId = await requireUserId(request); + const { organizationSlug, projectParam, envParam } = EnvironmentParamSchema.parse(params); + + const project = await findProjectBySlug(organizationSlug, projectParam, userId); + const redirectPath = concurrencyLimitsPath( + { slug: organizationSlug }, + { slug: projectParam }, + { slug: envParam } + ); + + if (!project) { + throw await redirectWithErrorMessage(redirectPath, request, "Project not found"); + } + + const formData = await request.formData(); + const submission = parseWithZod(formData, { schema: FormSchema }); + + if (submission.status !== "success") { + return json(submission.reply()); + } + + if (isPaidAddOnPurchase(submission.value.action)) { + const auth = await rbac.authenticateSession(request, { + userId, + organizationId: project.organizationId, + }); + if (!auth.ok || !auth.ability.can("manage", { type: "billing" })) { + return json( + submission.reply({ + fieldErrors: { amount: ["You don't have permission to manage billing."] }, + }), + { status: 403 } + ); + } + } + + if (submission.value.action === "allocate") { + const allocate = new AllocateConcurrencyService(); + const [error, result] = await tryCatch( + allocate.call({ + userId, + projectId: project.id, + organizationId: project.organizationId, + environments: submission.value.environments, + }) + ); + + if (error) { + return json( + submission.reply({ + fieldErrors: { + environments: [error instanceof Error ? error.message : "Unknown error"], + }, + }) + ); + } + + if (!result.success) { + return json(submission.reply({ fieldErrors: { environments: [result.error] } })); + } + + return redirectWithSuccessMessage( + `${redirectPath}?success=true`, + request, + "Concurrency allocated successfully" + ); + } + + const currentPlan = await getCurrentPlan(project.organizationId); + const purchaseBlockReason = getSelfServePurchaseBlockReason(currentPlan); + if (purchaseBlockReason === "plan_unavailable") { + return json( + submission.reply({ + fieldErrors: { amount: ["Unable to verify billing status. Please try again."] }, + }), + { status: 503 } + ); + } + if (purchaseBlockReason === "managed_billing") { + return json( + submission.reply({ fieldErrors: { amount: ["Contact us to request more concurrency."] } }), + { status: 403 } + ); + } + + const service = new SetConcurrencyAddOnService(); + const [error, result] = await tryCatch( + service.call({ + userId, + projectId: project.id, + organizationId: project.organizationId, + action: submission.value.action, + amount: submission.value.amount, + }) + ); + + if (error) { + return json( + submission.reply({ + fieldErrors: { amount: [error instanceof Error ? error.message : "Unknown error"] }, + }) + ); + } + + if (!result.success) { + return json(submission.reply({ fieldErrors: { amount: [result.error] } })); + } + + return redirectWithSuccessMessage( + `${redirectPath}?success=true`, + request, + submission.value.action === "purchase" + ? "Concurrency updated successfully" + : "Requested extra concurrency, we'll get back to you soon." + ); +}; + +export default function Page() { + const { + canAddConcurrency, + extraConcurrency, + extraAllocatedConcurrency, + extraUnallocatedConcurrency, + environments, + concurrencyPricing, + maxQuota, + canManageBilling, + } = useTypedLoaderData(); + + return ( + + + + + + + {environments.map((environment) => ( + + + {environment.type}{" "} + {environment.branchName ? ` (${environment.branchName})` : ""} + + + + + + ))} + + + + + + + {canAddConcurrency ? ( + + ) : ( + + )} + + + + ); +} + +function initialAllocation(environments: ConcurrencyResult["environments"]) { + return new Map( + environments + .filter((e) => e.type !== "DEVELOPMENT") + .map((e) => [e.id, Math.max(0, e.maximumConcurrencyLimit - e.planConcurrencyLimit)]) + ); +} + +function allocationTotal(environments: ConcurrencyResult["environments"]) { + const allocation = initialAllocation(environments); + return Array.from(allocation.values()).reduce((e, acc) => e + acc, 0); +} + +function Upgradable({ + extraConcurrency, + extraAllocatedConcurrency, + extraUnallocatedConcurrency, + environments, + concurrencyPricing, + maxQuota, + canManageBilling, +}: ConcurrencyResult & { canManageBilling: boolean }) { + const lastSubmission = useActionData(); + const [form, fields] = useForm({ + id: "allocate-concurrency", + // TODO: type this + lastResult: lastSubmission as any, + onValidate({ formData }) { + return parseWithZod(formData, { schema: FormSchema }); + }, + shouldRevalidate: "onSubmit", + }); + const { environments: formEnvironments } = fields; + + const navigation = useNavigation(); + const isLoading = navigation.state !== "idle" && navigation.formMethod === "POST"; + + const [allocation, setAllocation] = useState(initialAllocation(environments)); + + const allocatedInProject = Array.from(allocation.values()).reduce((e, acc) => e + acc, 0); + const initialAllocationInProject = allocationTotal(environments); + const changeInAllocation = allocatedInProject - initialAllocationInProject; + const unallocated = extraUnallocatedConcurrency - changeInAllocation; + const allocationModified = changeInAllocation !== 0; + + return ( +
+
+ Manage your concurrency +
+ + Concurrency limits determine how many runs you can execute at the same time. You can add + extra concurrency to your organization which you can allocate to environments in your + projects. + +
+
+
+ Extra concurrency + +
+ + + + Extra concurrency purchased + + {extraConcurrency} + + + + Allocated concurrency + + {allocationModified ? ( + <> + + {extraAllocatedConcurrency} + {" "} + {extraAllocatedConcurrency + changeInAllocation} + + ) : ( + extraAllocatedConcurrency + )} + + + + Unallocated concurrency + 0 + ? "text-success" + : unallocated < 0 + ? "text-error" + : "text-text-bright" + )} + > + {allocationModified ? ( + <> + + {extraUnallocatedConcurrency} + {" "} + {extraUnallocatedConcurrency - changeInAllocation} + + ) : ( + extraUnallocatedConcurrency + )} + + + 0 ? undefined : "after:bg-transparent" + } + > + 0 || allocationModified) && "pr-0")} + > +
+ {allocationModified ? ( + unallocated < 0 ? ( +
+ + + You're trying to allocate more concurrency than your total purchased + amount. + +
+ ) : ( +
+
+ + + Save your changes or{" "} + + . + +
+ +
+ ) + ) : unallocated > 0 ? ( +
+
+ + + You have {unallocated} extra concurrency available to allocate below. + +
+ +
+ ) : null} +
+ + + +
+ {formEnvironments.errors} +
+ + +
+ Concurrency allocation +
+ + + + Environment + + + Included{" "} + + + + Extra concurrency + Total + + + + {environments.map((environment, index) => ( + + + + + {environment.planConcurrencyLimit} + +
+ {environment.type === "DEVELOPMENT" ? ( + Math.max( + 0, + environment.maximumConcurrencyLimit - environment.planConcurrencyLimit + ) + ) : ( + <> + + { + const value = e.target.value === "" ? 0 : Number(e.target.value); + setAllocation(new Map(allocation).set(environment.id, value)); + }} + min={0} + /> + + )} +
+ + + {environment.type === "DEVELOPMENT" + ? environment.maximumConcurrencyLimit + : environment.planConcurrencyLimit + (allocation.get(environment.id) ?? 0)} + + + ))} + +
+ +
+
+ ); +} + +function NotUpgradable({ environments }: { environments: EnvironmentWithConcurrency[] }) { + const { isManagedCloud } = useFeatures(); + const plan = useCurrentPlan(); + const organization = useOrganization(); + const showSelfServe = useShowSelfServe(); + + return ( +
+
+ Your concurrency +
+ {isManagedCloud ? ( + <> + + Concurrency limits determine how many runs you can execute at the same time. You can + upgrade your plan to get more concurrency. You are currently on the{" "} + {plan?.v3Subscription?.plan?.title ?? "Free"} plan. + + {showSelfServe ? ( + + Upgrade for more concurrency + + ) : ( + Contact us} + /> + )} + + ) : null} +
+ + + + Environment + Concurrency limit + + + + {environments.map((environment) => ( + + + + + {environment.maximumConcurrencyLimit} + + ))} + +
+
+
+ ); +} + +function PurchaseConcurrencyModal({ + concurrencyPricing, + extraConcurrency, + extraUnallocatedConcurrency, + maxQuota, + canManageBilling, + disabled, +}: { + concurrencyPricing: { + stepSize: number; + centsPerStep: number; + }; + extraConcurrency: number; + extraUnallocatedConcurrency: number; + maxQuota: number; + canManageBilling: boolean; + disabled: boolean; +}) { + const showSelfServe = useShowSelfServe(); + const lastSubmission = useActionData(); + const [form, fields] = useForm({ + id: "purchase-concurrency", + // TODO: type this + lastResult: lastSubmission as any, + onValidate({ formData }) { + return parseWithZod(formData, { schema: FormSchema }); + }, + shouldRevalidate: "onSubmit", + }); + const { amount } = fields; + + const [amountValue, setAmountValue] = useState(extraConcurrency); + const navigation = useNavigation(); + const isLoading = navigation.state !== "idle" && navigation.formMethod === "POST"; + + // Close the panel, when we've succeeded + // This is required because a redirect to the same path doesn't clear state + const [searchParams, setSearchParams] = useSearchParams(); + const purchaseSucceeded = Boolean(searchParams.get("success")); + const [open, setOpen] = useState(false); + useEffect(() => { + if (purchaseSucceeded) { + // oxlint-disable-next-line react/set-state-in-effect -- This effect intentionally synchronizes route state after an external or lifecycle change. + setOpen(false); + setSearchParams((s) => { + s.delete("success"); + return s; + }); + } + }, [purchaseSucceeded, setSearchParams]); + + const state = updateState({ + value: amountValue, + existingValue: extraConcurrency, + quota: maxQuota, + extraUnallocatedConcurrency, + }); + const changeClassName = + state === "decrease" ? "text-error" : state === "increase" ? "text-success" : undefined; + + const title = extraConcurrency === 0 ? "Purchase extra concurrency" : "Add/remove concurrency"; + + if (!showSelfServe) { + return ( + Request more} + /> + ); + } + + return ( + + + + + + {title} +
+
+ + You can purchase bundles of {concurrencyPricing.stepSize} concurrency for{" "} + {formatCurrency(concurrencyPricing.centsPerStep / 100, false)}/month. Or you can + remove any extra concurrency after you have unallocated it from your environments + first. + +
+ + + setAmountValue(Number(e.target.value))} + disabled={isLoading} + /> + {amount.errors} + {form.errors} + +
+ {state === "need_to_increase_unallocated" ? ( +
+ + You need to unallocate{" "} + {formatNumber(extraConcurrency - amountValue - extraUnallocatedConcurrency)} more + concurrency from your environments in order to remove{" "} + {formatNumber(extraConcurrency - amountValue)} concurrency from your account. + +
+ ) : state === "above_quota" ? ( +
+ + Currently you can only have up to {maxQuota} extra concurrency. Send a request + below to lift your current limit. We'll get back to you soon. + +
+ ) : ( +
+
+ Summary + Total +
+
+ + {formatNumber(extraConcurrency)}{" "} + current total + + + {formatCurrency( + (extraConcurrency * concurrencyPricing.centsPerStep) / + concurrencyPricing.stepSize / + 100, + true + )} + +
+
+ + ({simplur`${extraConcurrency / concurrencyPricing.stepSize} bundle[|s]`}) + + /mth +
+
+ + {state === "increase" ? "+" : null} + {formatNumber(amountValue - extraConcurrency)} + + + {state === "increase" ? "+" : null} + {formatCurrency( + ((amountValue - extraConcurrency) * concurrencyPricing.centsPerStep) / + concurrencyPricing.stepSize / + 100, + true + )} + +
+
+ + ( + {simplur`${ + (amountValue - extraConcurrency) / concurrencyPricing.stepSize + } bundle[|s]`}{" "} + @ {formatCurrency(concurrencyPricing.centsPerStep / 100, true)}/mth) + + /mth +
+
+ + {formatNumber(amountValue)} new total + + + {formatCurrency( + (amountValue * concurrencyPricing.centsPerStep) / + concurrencyPricing.stepSize / + 100, + true + )} + +
+
+ + ({simplur`${amountValue / concurrencyPricing.stepSize} bundle[|s]`}) + + /mth +
+
+ )} +
+ + + + + ) : state === "decrease" || state === "need_to_increase_unallocated" ? ( + <> + + + + ) : ( + <> + + + + ) + } + cancelButton={ + + + + } + /> + +
+ + ); +} + +function updateState({ + value, + existingValue, + quota, + extraUnallocatedConcurrency, +}: { + value: number; + existingValue: number; + quota: number; + extraUnallocatedConcurrency: number; +}): "no_change" | "increase" | "decrease" | "above_quota" | "need_to_increase_unallocated" { + if (value === existingValue) return "no_change"; + if (value < existingValue) { + const difference = existingValue - value; + if (difference > extraUnallocatedConcurrency) { + return "need_to_increase_unallocated"; + } + return "decrease"; + } + if (value > quota) return "above_quota"; + return "increase"; +} diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.concurrency/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.concurrency/route.tsx index 1e5a015ec53..efd1e7cf5de 100644 --- a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.concurrency/route.tsx +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.concurrency/route.tsx @@ -1,93 +1,160 @@ -import { getFormProps, getInputProps, useForm } from "@conform-to/react"; -import { parseWithZod } from "@conform-to/zod/v4"; import { - ArrowDownIcon, - EnvelopeIcon, + ArrowUpCircleIcon, + BookOpenIcon, ExclamationTriangleIcon, - InformationCircleIcon, + PauseIcon, + PlayIcon, + RectangleStackIcon, } from "@heroicons/react/20/solid"; import { DialogClose } from "@radix-ui/react-dialog"; -import { Form, useActionData, useNavigation, useSearchParams } from "@remix-run/react"; -import { json, type ActionFunctionArgs, type LoaderFunctionArgs } from "@remix-run/server-runtime"; -import { tryCatch } from "@trigger.dev/core"; -import { useEffect, useState } from "react"; +import { Form, useNavigation } from "@remix-run/react"; +import { type LoaderFunctionArgs } from "@remix-run/server-runtime"; +import type { RuntimeEnvironmentType } from "@trigger.dev/database"; +import { useEffect, useMemo, useState, type ReactNode } from "react"; +import { QueuesIcon } from "~/assets/icons/QueuesIcon"; import { typedjson, useTypedLoaderData } from "remix-typedjson"; -import simplur from "simplur"; import { z } from "zod"; +import { ConcurrencyIcon } from "~/assets/icons/ConcurrencyIcon"; +import { RunsIcon } from "~/assets/icons/RunsIcon"; import { AdminDebugTooltip } from "~/components/admin/debugTooltip"; -import { CopyableText } from "~/components/primitives/CopyableText"; -import { EnvironmentCombo } from "~/components/environments/EnvironmentLabel"; -import { Feedback } from "~/components/Feedback"; -import { - MainHorizontallyCenteredContainer, - PageBody, - PageContainer, -} from "~/components/layout/AppLayout"; +import { environmentFullTitle } from "~/components/environments/EnvironmentLabel"; +import { PageBody, PageContainer } from "~/components/layout/AppLayout"; +import { MetricsLayout } from "~/components/layout/MetricsLayout"; +import { Badge } from "~/components/primitives/Badge"; import { Button, LinkButton } from "~/components/primitives/Buttons"; import { Dialog, DialogContent, DialogHeader, DialogTrigger } from "~/components/primitives/Dialog"; -import { Fieldset } from "~/components/primitives/Fieldset"; import { FormButtons } from "~/components/primitives/FormButtons"; -import { FormError } from "~/components/primitives/FormError"; -import { Header2, Header3 } from "~/components/primitives/Headers"; -import { Input } from "~/components/primitives/Input"; -import { InputGroup } from "~/components/primitives/InputGroup"; -import { InputNumberStepper } from "~/components/primitives/InputNumberStepper"; -import { Label } from "~/components/primitives/Label"; +import { Header3 } from "~/components/primitives/Headers"; import { NavBar, PageAccessories, PageTitle } from "~/components/primitives/PageHeader"; +import { PaginationControls } from "~/components/primitives/Pagination"; import { Paragraph } from "~/components/primitives/Paragraph"; -import * as Property from "~/components/primitives/PropertyTable"; -import { SpinnerWhite } from "~/components/primitives/Spinner"; +import { PopoverMenuItem } from "~/components/primitives/Popover"; +import { SearchInput } from "~/components/primitives/SearchInput"; +import { Spinner } from "~/components/primitives/Spinner"; import { Table, TableBody, TableCell, + TableCellMenu, TableHeader, TableHeaderCell, TableRow, } from "~/components/primitives/Table"; -import { InfoIconTooltip } from "~/components/primitives/Tooltip"; -import { useFeatures } from "~/hooks/useFeatures"; +import { + InfoIconTooltip, + SimpleTooltip, + Tooltip, + TooltipContent, + TooltipProvider, + TooltipTrigger, +} from "~/components/primitives/Tooltip"; +import { TasksIcon } from "~/assets/icons/TasksIcon"; +import { QueueName } from "~/components/runs/v3/QueueName"; +import { env } from "~/env.server"; +import { useAutoRevalidate } from "~/hooks/useAutoRevalidate"; +import { useEnvironment } from "~/hooks/useEnvironment"; import { useOrganization } from "~/hooks/useOrganizations"; -import { useShowSelfServe } from "~/hooks/useShowSelfServe"; +import { useProject } from "~/hooks/useProject"; import { redirectWithErrorMessage, redirectWithSuccessMessage } from "~/models/message.server"; import { findProjectBySlug } from "~/models/project.server"; +import { findEnvironmentBySlug } from "~/models/runtimeEnvironment.server"; +import { EnvironmentQueuePresenter } from "~/presenters/v3/EnvironmentQueuePresenter.server"; +import { QueueListPresenter } from "~/presenters/v3/QueueListPresenter.server"; import { - ManageConcurrencyPresenter, - type ConcurrencyResult, - type EnvironmentWithConcurrency, -} from "~/presenters/v3/ManageConcurrencyPresenter.server"; + QueueMetricsPresenter, + type QueueListMetric, +} from "~/presenters/v3/QueueMetricsPresenter.server"; +import { TimeFilter, timeFilterFromTo } from "~/components/runs/v3/SharedFilters"; +import { useSearchParams } from "~/hooks/useSearchParam"; +import { parseFiniteInt } from "~/utils/searchParams"; +import { MiniLineChart } from "~/components/metrics/MiniLineChart"; +import { buildActivityTimeAxis } from "~/components/primitives/charts/activityTimeAxis"; +import { Chart, type ChartConfig } from "~/components/primitives/charts/ChartCompound"; +import { ChartCard } from "~/components/primitives/charts/ChartCard"; +import { ChartSyncProvider } from "~/components/primitives/charts/ChartSyncContext"; +import { useZoomToTimeFilter } from "~/hooks/useZoomToTimeFilter"; import { - getCurrentPlan, - getPlans, - getSelfServePurchaseBlockReason, -} from "~/services/platform.v3.server"; -import { textLinkClassName } from "~/components/primitives/TextLink"; + useIsMetricResponseFresh, + useMetricResourceQuery, + type MetricResourceTimeRange, +} from "~/hooks/useMetricResourceQuery"; +import { logger } from "~/services/logger.server"; import { rbac } from "~/services/rbac.server"; +import { resolveProjectAuthScope } from "~/services/projectAuthScope.server"; +import { dashboardAction } from "~/services/routeBuilders/dashboardBuilder"; import { requireUserId } from "~/services/session.server"; import { cn } from "~/utils/cn"; -import { formatCurrency, formatNumber } from "~/utils/numberFormatter"; -import { isPaidAddOnPurchase } from "~/utils/paidAddOnPermissions"; -import { concurrencyPath, EnvironmentParamSchema, v3BillingPath } from "~/utils/pathBuilder"; -import { AllocateConcurrencyService } from "~/v3/services/allocateConcurrency.server"; -import { SetConcurrencyAddOnService } from "~/v3/services/setConcurrencyAddOn.server"; -import { useCurrentPlan } from "../_app.orgs.$organizationSlug/route"; -import { sectionAgentPageContext } from "~/components/dashboard-agent/suggested-prompts"; +import { ENVIRONMENT_PAUSE_SOURCE_BILLING_LIMIT } from "~/utils/environmentPauseSource"; +import { + concurrencyLimitsPath, + docsPath, + EnvironmentParamSchema, + v3BillingPath, + concurrencyQueuePath, + v3RunsPath, +} from "~/utils/pathBuilder"; import type { Handle } from "~/utils/handle"; +import { queuesAgentPageContext } from "~/components/dashboard-agent/suggested-prompts"; +import { WhenAgentUnavailable } from "~/components/dashboard-agent/WhenAgentUnavailable"; +import { PauseEnvironmentService } from "~/v3/services/pauseEnvironment.server"; +import { handleQueueMutationAction } from "~/models/queueMutation.server"; +import { + QueueOverrideConcurrencyButton, + QueuePauseResumeButton, +} from "~/components/queues/QueueControls"; +import { useCurrentPlan } from "../_app.orgs.$organizationSlug/route"; +import { BigNumber } from "~/components/metrics/BigNumber"; +import { canAccessQueueMetricsUi } from "~/v3/canAccessQueueMetricsUi.server"; +import { QueueAllocationPresenter } from "~/presenters/v3/QueueAllocationPresenter.server"; +import { + QUEUE_METRICS_DEFAULT_PERIOD, + QUEUE_METRICS_RETENTION_DAYS, + clampQueueMetricsPeriod, + clipQueueMetricsWindow, + queueMetricsPeriodFromRequest, + resolveQueueMetricsPeriod, + useRememberQueueMetricsPeriod, +} from "~/components/queues/queueMetricsPeriod"; +import { queueMetricsMaxPeriodDays } from "~/components/queues/queueMetricsPeriod.server"; +import { isQueueAtCapacity } from "~/components/queues/queue-thresholds"; +import { pageMeta } from "~/utils/pageTitle"; +import { InlineCode } from "~/components/code/InlineCode"; + +const SearchParamsSchema = z.object({ + query: z.string().optional(), + page: z.coerce.number().min(1).default(1), + period: z.string().optional(), + from: z.string().optional(), + to: z.string().optional(), + sort: z.enum(["busiest", "queued", "name"]).optional(), +}); + +// The live "Queued" / "Running" header blocks poll ClickHouse on a short cadence so they stay +// current after first paint. They read the env-wide gauges from env_metrics (the env-level rollup +// of concurrency_metrics, cheapest for a dimension-free query), always over a fixed 15m window regardless +// of the chart/table period, and are NOT scoped to the visible queue set (the blocks are env-wide). +const QUEUE_LIVE_BLOCKS_PERIOD = "15m"; +const QUEUE_LIVE_BLOCKS_QUERY = + "SELECT timeBucket() AS t, max(max_env_queued) AS env_queued, max(max_env_running) AS env_running FROM env_metrics GROUP BY t ORDER BY t"; +// Trust the ClickHouse gauge only while its newest bucket is this recent; otherwise fall back to +// the loader's Redis-exact live values (matches LIVE_GAUGE_FRESH_MS on the queue detail page / run +// inspector). +const LIVE_GAUGE_FRESH_MS = 90_000; export const handle: Handle = { - agentPageContext: () => sectionAgentPageContext("concurrency"), + agentPageContext: (data) => queuesAgentPageContext(data), }; -import { pageMeta } from "~/utils/pageTitle"; -export const meta = pageMeta("Manage concurrency"); +export const meta = pageMeta("Concurrency"); export const loader = async ({ request, params }: LoaderFunctionArgs) => { const userId = await requireUserId(request); - const { - organizationSlug, - projectParam, - envParam: _envParam, - } = EnvironmentParamSchema.parse(params); + const { organizationSlug, projectParam, envParam } = EnvironmentParamSchema.parse(params); + + const url = new URL(request.url); + const { page, query, period, from, to, sort } = SearchParamsSchema.parse( + Object.fromEntries(url.searchParams) + ); const project = await findProjectBySlug(organizationSlug, projectParam, userId); if (!project) { @@ -97,827 +164,2130 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { }); } + const environment = await findEnvironmentBySlug(project.id, envParam, userId); + if (!environment) { + throw new Response(undefined, { + status: 404, + statusText: "Environment not found", + }); + } + const auth = await rbac.authenticateSession(request, { userId, organizationId: project.organizationId, + projectId: project.id, }); - const canManageBilling = auth.ok && auth.ability.can("manage", { type: "billing" }); + const canWriteTasks = auth.ok && auth.ability.can("write", { type: "tasks" }); - const presenter = new ManageConcurrencyPresenter(); - const [error, result] = await tryCatch( - presenter.call({ - userId: userId, - projectId: project.id, - organizationId: project.organizationId, - }) + // Per-org gate for the metrics UI. When off, this org gets the classic Queues page and + // no metrics query fires. + const queueMetricsUiEnabled = await canAccessQueueMetricsUi({ + request, + userId, + organizationSlug, + }); + + const maxPeriodDays = queueMetricsUiEnabled + ? await queueMetricsMaxPeriodDays(environment.organizationId) + : QUEUE_METRICS_RETENTION_DAYS; + const defaultPeriod = clampQueueMetricsPeriod( + queueMetricsPeriodFromRequest(request), + maxPeriodDays ); - if (error) { - throw new Response(undefined, { - status: 400, - statusText: error.message, + try { + const queueListPresenter = new QueueListPresenter(); + const queues = await queueListPresenter.call({ + environment, + query, + page, + includeLimits: true, + // Relevance ordering rides the metrics pipeline, so it is part of the gated UI. + sort: queueMetricsUiEnabled ? (sort ?? "busiest") : "name", }); - } - const plans = await tryCatch(getPlans()); - if (!plans) { - throw new Response(null, { status: 404, statusText: "Plans not found" }); - } + const environmentQueuePresenter = new EnvironmentQueuePresenter(); - return typedjson({ ...result, canManageBilling }); -}; + const autoReloadPollIntervalMs = env.QUEUES_AUTORELOAD_POLL_INTERVAL_MS; -const FormSchema = z.discriminatedUnion("action", [ - z.object({ - action: z.enum(["purchase"]), - amount: z.coerce.number().min(0, "Amount must be 0 or more"), - }), - z.object({ - action: z.enum(["quota-increase"]), - amount: z.coerce.number().min(1, "Amount must be greater than 0"), - }), - z.object({ - action: z.enum(["allocate"]), - // It will only update environments that are passed in - environments: z.array( - z.object({ - id: z.string(), - amount: z.coerce.number().min(0, "Amount must be 0 or more"), - }) - ), - }), -]); - -export const action = async ({ request, params }: ActionFunctionArgs) => { - const userId = await requireUserId(request); - const { organizationSlug, projectParam, envParam } = EnvironmentParamSchema.parse(params); + // Per-queue list metrics (Delay p95 + backlog sparkline columns) are SSR'd with the table. + // The environment header tiles are fetched client-side per card (see QueueEnvMetricChart) so a + // slow ClickHouse query never blocks the queues list from rendering. + let metrics: { + bucketStartMs: number; + bucketIntervalMs: number; + byQueue: Record; + } | null = null; - const project = await findProjectBySlug(organizationSlug, projectParam, userId); - const redirectPath = concurrencyPath( - { slug: organizationSlug }, - { slug: projectParam }, - { slug: envParam } - ); - - if (!project) { - throw await redirectWithErrorMessage(redirectPath, request, "Project not found"); - } + if (queueMetricsUiEnabled) { + // Metrics are additive observability; a ClickHouse hiccup must not take down queue + // management. Fail open to metrics: null instead of bubbling to the page-level 400. + try { + const presenter = new QueueMetricsPresenter(); + const queueNames = queues.queues.map((q) => + q.type === "task" ? `task/${q.name}` : q.name + ); + const timeRange = clipQueueMetricsWindow( + timeFilterFromTo({ + period: + resolveQueueMetricsPeriod({ + period, + from, + to, + defaultPeriod, + maxPeriodDays, + }) ?? undefined, + from: parseFiniteInt(from), + to: parseFiniteInt(to), + defaultPeriod, + }), + maxPeriodDays + ); + const queueMetrics = + queueNames.length > 0 + ? await presenter.getQueueListMetrics({ + environment, + queueNames, + from: timeRange.from, + to: timeRange.to, + }) + : null; + if (queueMetrics) { + metrics = { + bucketStartMs: queueMetrics.bucketStartMs, + bucketIntervalMs: queueMetrics.bucketIntervalMs, + byQueue: Object.fromEntries(queueMetrics.byQueue), + }; + } + } catch (error) { + logger.warn("Queue list metrics unavailable, rendering without them", { + error, + organizationSlug: params.organizationSlug, + }); + } + } - const formData = await request.formData(); - const submission = parseWithZod(formData, { schema: FormSchema }); + // Allocation summary (Environment limit + Allocated tiles) is additive; a presenter + // failure must not 400 the page, so fail open to null like the metrics block above. + let allocation: Awaited> | null = null; + if (queueMetricsUiEnabled) { + try { + allocation = await new QueueAllocationPresenter().call({ environment }); + } catch (error) { + logger.warn("Queue allocation summary unavailable, rendering without it", { + error, + organizationSlug: params.organizationSlug, + }); + } + } - if (submission.status !== "success") { - return json(submission.reply()); + return typedjson({ + ...queues, + environment: await environmentQueuePresenter.call(environment), + autoReloadPollIntervalMs, + metrics, + allocation, + queueMetricsUiEnabled, + defaultPeriod, + maxPeriodDays, + canWriteTasks, + }); + } catch (error) { + console.error(error); + throw new Response(undefined, { + status: 400, + statusText: "Something went wrong, if this problem persists please contact support.", + }); } +}; - if (isPaidAddOnPurchase(submission.value.action)) { - const auth = await rbac.authenticateSession(request, { - userId, - organizationId: project.organizationId, - }); - if (!auth.ok || !auth.ability.can("manage", { type: "billing" })) { - return json( - submission.reply({ - fieldErrors: { amount: ["You don't have permission to manage billing."] }, - }), - { status: 403 } +export const action = dashboardAction( + { + params: EnvironmentParamSchema, + context: (params) => resolveProjectAuthScope(params.organizationSlug, params.projectParam), + authorization: { + action: "write", + resource: { type: "tasks" }, + message: "With your current role, you can't manage queues.", + }, + }, + async ({ request, params, user }) => { + const userId = user.id; + if (request.method.toLowerCase() !== "post") { + return redirectWithErrorMessage( + `/orgs/${params.organizationSlug}/projects/${params.projectParam}/env/${params.envParam}/concurrency`, + request, + "Wrong method" ); } - } - if (submission.value.action === "allocate") { - const allocate = new AllocateConcurrencyService(); - const [error, result] = await tryCatch( - allocate.call({ - userId, - projectId: project.id, - organizationId: project.organizationId, - environments: submission.value.environments, - }) - ); - - if (error) { - return json( - submission.reply({ - fieldErrors: { - environments: [error instanceof Error ? error.message : "Unknown error"], - }, - }) - ); + const { organizationSlug, projectParam, envParam } = EnvironmentParamSchema.parse(params); + + const project = await findProjectBySlug(organizationSlug, projectParam, userId); + if (!project) { + throw new Response(undefined, { + status: 404, + statusText: "Project not found", + }); } - if (!result.success) { - return json(submission.reply({ fieldErrors: { environments: [result.error] } })); + const environment = await findEnvironmentBySlug(project.id, envParam, userId); + if (!environment) { + throw new Response(undefined, { + status: 404, + statusText: "Environment not found", + }); } - return redirectWithSuccessMessage( - `${redirectPath}?success=true`, - request, - "Concurrency allocated successfully" - ); - } + const formData = await request.formData(); + const action = formData.get("action"); - const currentPlan = await getCurrentPlan(project.organizationId); - const purchaseBlockReason = getSelfServePurchaseBlockReason(currentPlan); - if (purchaseBlockReason === "plan_unavailable") { - return json( - submission.reply({ - fieldErrors: { amount: ["Unable to verify billing status. Please try again."] }, - }), - { status: 503 } - ); - } - if (purchaseBlockReason === "managed_billing") { - return json( - submission.reply({ fieldErrors: { amount: ["Contact us to request more concurrency."] } }), - { status: 403 } - ); - } + const url = new URL(request.url); + const redirectPath = `/orgs/${organizationSlug}/projects/${projectParam}/env/${envParam}/concurrency${url.search}`; - const service = new SetConcurrencyAddOnService(); - const [error, result] = await tryCatch( - service.call({ + if (environment.archivedAt) { + return redirectWithErrorMessage(redirectPath, request, "This branch is archived"); + } + + // Per-queue actions (pause/resume/override/remove-override) are shared with the queue detail + // route, so they live in a helper that both routes call. + const queueMutation = await handleQueueMutationAction({ + request, + environment, userId, - projectId: project.id, - organizationId: project.organizationId, - action: submission.value.action, - amount: submission.value.amount, - }) - ); + formData, + redirectPath, + }); + if (queueMutation) { + return queueMutation; + } - if (error) { - return json( - submission.reply({ - fieldErrors: { amount: [error instanceof Error ? error.message : "Unknown error"] }, - }) - ); + switch (action) { + case "environment-pause": { + const pauseService = new PauseEnvironmentService(); + const result = await pauseService.call(environment, "paused"); + if (!result.success) { + return redirectWithErrorMessage(redirectPath, request, result.error); + } + return redirectWithSuccessMessage(redirectPath, request, "Environment paused"); + } + case "environment-resume": { + const resumeService = new PauseEnvironmentService(); + const result = await resumeService.call(environment, "resumed"); + if (!result.success) { + return redirectWithErrorMessage(redirectPath, request, result.error); + } + return redirectWithSuccessMessage(redirectPath, request, "Environment resumed"); + } + default: + return redirectWithErrorMessage(redirectPath, request, "Something went wrong"); + } } +); - if (!result.success) { - return json(submission.reply({ fieldErrors: { amount: [result.error] } })); - } +// Derives the environment concurrency status ("limit" | "burst" | "within") and the matching +// text color from the current running count vs. the env limit and burst factor. Shared by both +// the classic and metrics views so the "Running" tile styling stays in sync. +function getEnvConcurrencyLimitStatus(environment: { + running: number; + concurrencyLimit: number; + burstFactor: number; +}) { + const limitStatus = + environment.running === environment.concurrencyLimit * environment.burstFactor + ? "limit" + : environment.running > environment.concurrencyLimit + ? "burst" + : "within"; - return redirectWithSuccessMessage( - `${redirectPath}?success=true`, - request, - submission.value.action === "purchase" - ? "Concurrency updated successfully" - : "Requested extra concurrency, we'll get back to you soon." - ); -}; + const limitClassName = + limitStatus === "burst" ? "text-warning" : limitStatus === "limit" ? "text-error" : undefined; + + return { limitStatus, limitClassName }; +} export default function Page() { + // Per-org flag decides which whole page renders. Off => the classic Queues page, + // byte-for-byte the pre-metrics UI. Each branch is its own component (own hooks). + const { queueMetricsUiEnabled } = useTypedLoaderData(); + return queueMetricsUiEnabled ? : ; +} + +function QueuesWithMetricsView() { const { - canAddConcurrency, - extraConcurrency, - extraAllocatedConcurrency, - extraUnallocatedConcurrency, - environments, - concurrencyPricing, - maxQuota, - canManageBilling, + environment, + queues, + pagination, + totalQueues, + hasFilters, + autoReloadPollIntervalMs, + metrics, + allocation, + defaultPeriod, + maxPeriodDays, + canWriteTasks, } = useTypedLoaderData(); + const metricsByQueue = metrics?.byQueue ?? {}; + + const organization = useOrganization(); + const project = useProject(); + const env = useEnvironment(); + const plan = useCurrentPlan(); + + // The header tiles fetch client-side with the same period/from/to the TimeFilter writes. + const { value } = useSearchParams(); + const timeRange = { + period: resolveQueueMetricsPeriod({ + period: value("period"), + from: value("from"), + to: value("to"), + defaultPeriod, + maxPeriodDays, + }), + from: value("from") ?? null, + to: value("to") ?? null, + }; + useRememberQueueMetricsPeriod(value("period")); + + useAutoRevalidate({ interval: autoReloadPollIntervalMs, onFocus: true }); + + // Drag-to-zoom on either chart narrows the page's from/to search params, which the + // TimeFilter and the client-side metric queries both read (same wiring as the Agent page). + const zoomToTimeFilter = useZoomToTimeFilter(); + + // Live env-wide Queued/Running blocks. First paint uses the loader's Redis-exact values; from the + // first poll on we prefer ClickHouse so the blocks stay current without a full page revalidate. + // Empty rows (quiet env, or the very first fetch still in flight) fall back to the loader values, + // so we never flash a stale 0. Fixed 15m window, env-wide (no queue filter), CH-only recurring + // load; pauses while the tab is hidden (handled inside the hook). + const { rows: liveBlockRows, responseReceivedAt } = useMetricResourceQuery( + QUEUE_LIVE_BLOCKS_QUERY, + { + organizationId: organization.id, + projectId: project.id, + environmentId: env.id, + timeRange: { period: QUEUE_LIVE_BLOCKS_PERIOD, from: null, to: null }, + defaultPeriod: QUEUE_LIVE_BLOCKS_PERIOD, + fillGaps: false, + refreshIntervalMs: 15_000, + } + ); + const lastLiveBlockRow = + liveBlockRows.length > 0 ? liveBlockRows[liveBlockRows.length - 1] : null; + // Only trust the gauge while its newest bucket is fresh. A row painted from the hook's cache on + // client-side nav-back (responseCache), or a quiet env whose latest bucket is minutes old, must + // not override the loader's Redis-exact live values with a stale count. + const lastLiveBucketMs = lastLiveBlockRow ? tileTimeToMs(lastLiveBlockRow.t) : NaN; + const liveBlockIsFresh = useIsMetricResponseFresh( + responseReceivedAt, + lastLiveBucketMs, + LIVE_GAUGE_FRESH_MS + ); + const freshLiveBlockRow = lastLiveBlockRow && liveBlockIsFresh ? lastLiveBlockRow : null; + const envQueuedLive = freshLiveBlockRow + ? tileNumber(freshLiveBlockRow.env_queued) + : environment.queued; + const envRunningLive = freshLiveBlockRow + ? tileNumber(freshLiveBlockRow.env_running) + : environment.running; + + // Allocation summary tiles. The presenter computes the env-wide allocated total (sum of + // each queue's explicit limit clamped to the env limit) in a single aggregate query. + const envLimit = environment.concurrencyLimit; + const burstLimit = Math.round(envLimit * environment.burstFactor); + const allocated = allocation?.allocated ?? 0; + const allocationPct = envLimit > 0 ? Math.round((allocated / envLimit) * 100) : 0; + + // Running-block tinting (burst/limit) tracks the live running value, not the loader snapshot. + const { limitStatus, limitClassName } = getEnvConcurrencyLimitStatus({ + running: envRunningLive, + concurrencyLimit: environment.concurrencyLimit, + burstFactor: environment.burstFactor, + }); + + // Client-side, header-click sorting over the current page's rows. Server pagination and the + // default busiest order are unchanged; clearing a sort returns to that server order. + const queueRows = queues ?? []; + return ( - - - {environments.map((environment) => ( - - - {environment.type}{" "} - {environment.branchName ? ` (${environment.branchName})` : ""} - - - - - - ))} - - + + + + Queues docs + + - - - {canAddConcurrency ? ( - + {/* Filters — pinned bar directly under the NavBar. This row is page-wide only: Period is + the one control that changes the tiles and charts below, so it leads the row. Search + and pagination scope the table alone and live in that table's own bar instead. */} + +
+ - ) : ( - - )} - - - - ); -} - -function initialAllocation(environments: ConcurrencyResult["environments"]) { - return new Map( - environments - .filter((e) => e.type !== "DEVELOPMENT") - .map((e) => [e.id, Math.max(0, e.maximumConcurrencyLimit - e.planConcurrencyLimit)]) - ); -} - -function allocationTotal(environments: ConcurrencyResult["environments"]) { - const allocation = initialAllocation(environments); - return Array.from(allocation.values()).reduce((e, acc) => e + acc, 0); -} - -function Upgradable({ - extraConcurrency, - extraAllocatedConcurrency, - extraUnallocatedConcurrency, - environments, - concurrencyPricing, - maxQuota, - canManageBilling, -}: ConcurrencyResult & { canManageBilling: boolean }) { - const lastSubmission = useActionData(); - const [form, fields] = useForm({ - id: "allocate-concurrency", - // TODO: type this - lastResult: lastSubmission as any, - onValidate({ formData }) { - return parseWithZod(formData, { schema: FormSchema }); - }, - shouldRevalidate: "onSubmit", - }); - const { environments: formEnvironments } = fields; - - const navigation = useNavigation(); - const isLoading = navigation.state !== "idle" && navigation.formMethod === "POST"; +
+
+ {environment.runsEnabled && + env.pauseSource !== ENVIRONMENT_PAUSE_SOURCE_BILLING_LIMIT ? ( + + ) : null} +
+ - const [allocation, setAllocation] = useState(initialAllocation(environments)); + {/* Queued + Running + Allocated + Environment limit summary. Four stat tiles: the grid + derives its columns from the tile count (two-up, four-up from lg). The allocation + presenter fails open to null (a ClickHouse/PG hiccup mustn't take down the tiles), so + only the Allocated tile depends on it — the other three + controls always render, and + Allocated shows a "–" placeholder to keep the 4-tile grid shape stable. */} + + paused : undefined} + animate + accessory={ + + + + } + valueClassName={env.paused ? "text-warning tabular-nums" : "tabular-nums"} + compactThreshold={1000000} + /> + + Including {envRunningLive - environment.concurrencyLimit} burst runs{" "} + + + ) : limitStatus === "limit" ? ( + "At concurrency limit" + ) : undefined + } + accessory={ + + + + } + compactThreshold={1000000} + /> + + Allocated + {allocation ? ( + + ) : null} + + } + value={allocation ? allocated : undefined} + formattedValue={allocation ? undefined : "–"} + suffix={allocation ? `${allocationPct}% of the environment limit` : undefined} + suffixClassName="text-text-dimmed" + /> + 1 ? `bursts up to ${burstLimit}` : undefined} + suffixClassName="text-text-dimmed" + accessory={ + plan ? ( + plan?.v3Subscription?.plan?.limits.concurrentRuns.canExceed ? ( + + Increase limit + + ) : ( + + Increase limit + + ) + ) : undefined + } + /> + - const allocatedInProject = Array.from(allocation.values()).reduce((e, acc) => e + acc, 0); - const initialAllocationInProject = allocationTotal(environments); - const changeInAllocation = allocatedInProject - initialAllocationInProject; - const unallocated = extraUnallocatedConcurrency - changeInAllocation; - const allocationModified = changeInAllocation !== 0; + {hasFilters || totalQueues !== 0 ? ( + + + {QUEUE_HEADER_TILES.map((tile) => ( + 1 + ? [ + { + y: Math.round(environment.burstFactor * 100), + label: `Burst ${Math.round( + environment.concurrencyLimit * environment.burstFactor + )}`, + labelPlacement: "outside" as const, + }, + ] + : []), + ] + : undefined + } + // Saturation recolours the line above its 100% limit with a gradient split, so + // only the portion over the line is orange (the offset is derived from the line's + // own value range, so the split lands exactly at 100% regardless of domain + // padding). p95 and throttled use a per-bucket overlay: it retraces only the + // over-threshold stretches, so under-threshold buckets stay blue. + thresholdStroke={ + tile.id === "saturation" + ? { value: 100, aboveColor: "var(--color-warning)" } + : undefined + } + warningOverlay={ + tile.id === "p95" + ? { threshold: 60_000 } + : tile.id === "throttled" + ? // Integer counts: threshold 0 warns once a bucket has ≥1 throttle. + { threshold: 0 } + : undefined + } + /> + ))} + + + ) : null} - return ( -
-
- Manage your concurrency -
- - Concurrency limits determine how many runs you can execute at the same time. You can add - extra concurrency to your organization which you can allocate to environments in your - projects. - -
-
-
- Extra concurrency - -
- - - - Extra concurrency purchased - - {extraConcurrency} - - - - Allocated concurrency - - {allocationModified ? ( - <> - - {extraAllocatedConcurrency} - {" "} - {extraAllocatedConcurrency + changeInAllocation} - - ) : ( - extraAllocatedConcurrency - )} - - + + + + + } + > + {/* Default overflow-x-auto container so wide tables still scroll horizontally on + narrow viewports; the page (not this region) owns vertical scrolling. */} +
+ - Unallocated concurrency - Name + Queued + Running + 0 - ? "text-success" - : unallocated < 0 - ? "text-error" - : "text-text-bright" - )} + disableTooltipHoverableContent + tooltip={limitTooltip} + tooltipContentClassName="max-w-xs" + > + Limit + + +

+ Environment: uses the environment + limit of {environment.concurrencyLimit}. +

+

+ User: a limit you set in your + code. +

+

+ Override: a limit you set here or + via the API. +

+ + } > - {allocationModified ? ( + Limited by + + Health + + Delay p95 + + - - {extraUnallocatedConcurrency} - {" "} - {extraUnallocatedConcurrency - changeInAllocation} + How many runs were waiting, over the selected time. marks + where the queue was throttled. - ) : ( - extraUnallocatedConcurrency - )} - - - 0 ? undefined : "after:bg-transparent" - } - > - 0 || allocationModified) && "pr-0")} + } > -
- {allocationModified ? ( - unallocated < 0 ? ( -
- - - You're trying to allocate more concurrency than your total purchased - amount. - -
- ) : ( -
-
- - - Save your changes or{" "} - - . - -
- -
- ) - ) : unallocated > 0 ? ( -
-
- - - You have {unallocated} extra concurrency available to allocate below. - -
- -
- ) : null} -
- - - -
- {formEnvironments.errors} -
-
- -
- Concurrency allocation -
- - - - Environment - - - Included{" "} - - + Backlog + + + Pause/resume - Extra concurrency - Total - {environments.map((environment, index) => ( - - - - - {environment.planConcurrencyLimit} - -
- {environment.type === "DEVELOPMENT" ? ( - Math.max( - 0, - environment.maximumConcurrencyLimit - environment.planConcurrencyLimit - ) - ) : ( - <> - - { - const value = e.target.value === "" ? 0 : Number(e.target.value); - setAllocation(new Map(allocation).set(environment.id, value)); - }} - min={0} + {queueRows.length > 0 ? ( + queueRows.map((queue) => { + const limit = queue.limits.perKey.current ?? environment.concurrencyLimit; + const isLimit = queue.kind === "limit"; + /** A limit row's `running` counts holders across every key, so only its + * total bound compares against it; a perKey-only limit has no aggregate + * threshold and never reads as at-limit here. */ + const atLimitThreshold = isLimit + ? queue.limits.total?.current != null + ? Math.min(queue.limits.total.current, environment.concurrencyLimit) + : null + : limit; + /** A zero threshold is a pause (nothing may run), not saturation — + * without the guard `running >= 0` holds for every row. */ + const isAtConcurrencyLimit = + atLimitThreshold !== null && + atLimitThreshold > 0 && + queue.running >= atLimitThreshold; + const isAtQueueLimit = + environment.queueSizeLimit !== null && + queue.queued >= environment.queueSizeLimit; + const queueFilterableName = queueMetricsKey(queue); + const queueMetric = metricsByQueue[queueFilterableName]; + /** The detail page is queue observability (queue metrics, run filters, pause + * and override actions); a limit's activity lives on each holder's home queue, + * so limit rows don't link anywhere. */ + const queueDetailPath = isLimit + ? undefined + : concurrencyQueuePath(organization, project, env, { + friendlyId: queue.id, + }); + const displayName = isLimit ? queue.name.replace(/^limit\//, "") : queue.name; + return ( + + s, so + // they render beside the link (leading/trailing), never inside it — + // otherwise the cell is invalid
- -
-
+ + + ); } -function NotUpgradable({ environments }: { environments: EnvironmentWithConcurrency[] }) { - const { isManagedCloud } = useFeatures(); - const plan = useCurrentPlan(); - const organization = useOrganization(); - const showSelfServe = useShowSelfServe(); +function EnvironmentPauseResumeButton({ + env, + disabled, +}: { + env: { type: RuntimeEnvironmentType; paused: boolean }; + disabled: boolean; +}) { + const navigation = useNavigation(); + const [isOpen, setIsOpen] = useState(false); + + useEffect(() => { + if (navigation.state === "loading" || navigation.state === "idle") { + // oxlint-disable-next-line react/set-state-in-effect -- This effect intentionally synchronizes route state after an external or lifecycle change. + setIsOpen(false); + } + }, [navigation.state]); + + const isLoading = Boolean( + navigation.formData?.get("action") === (env.paused ? "environment-resume" : "environment-pause") + ); return ( -
-
- Your concurrency + +
+ + + +
+ + + +
+ + + {disabled + ? "You don't have permission to manage queues" + : env.paused + ? `Resumes ${environmentFullTitle(env)} so its runs can be dequeued again.` + : `Pauses all runs from being dequeued in ${environmentFullTitle(env)}. Any executing runs will continue to run.`} + + +
- {isManagedCloud ? ( - <> - - Concurrency limits determine how many runs you can execute at the same time. You can - upgrade your plan to get more concurrency. You are currently on the{" "} - {plan?.v3Subscription?.plan?.title ?? "Free"} plan. + + {env.paused ? "Resume environment?" : "Pause environment?"} +
+ + {env.paused + ? `This will allow runs to be dequeued in ${environmentFullTitle(env)} again.` + : `This will pause all runs from being dequeued in ${environmentFullTitle( + env + )}. Any executing runs will continue to run.`} - {showSelfServe ? ( - - Upgrade for more concurrency - - ) : ( - Contact us} +
setIsOpen(false)}> + - )} - - ) : null} -
- - - - Environment - Concurrency limit - - - - {environments.map((environment) => ( - - - - - {environment.maximumConcurrencyLimit} - - ))} - -
-
-
+ : env.paused ? PlayIcon : PauseIcon + } + shortcut={{ modifiers: ["mod"], key: "enter" }} + > + {env.paused ? "Resume environment" : "Pause environment"} + + } + cancelButton={ + + + + } + /> + +
+ +
); } -function PurchaseConcurrencyModal({ - concurrencyPricing, - extraConcurrency, - extraUnallocatedConcurrency, - maxQuota, - canManageBilling, - disabled, -}: { - concurrencyPricing: { - stepSize: number; - centsPerStep: number; +export function isEnvironmentPauseResumeFormSubmission( + formMethod: string | undefined, + formData: FormData | undefined +) { + if (!formMethod || !formData) { + return false; + } + + return ( + formMethod.toLowerCase() === "post" && + (formData.get("action") === "environment-pause" || + formData.get("action") === "environment-resume") + ); +} + +export function QueueFilters() { + return ; +} + +type MetricTileRow = Record; + +type TilePoint = { bucket: number; value: number | null }; + +// Inline colour swatch matching the chart's warning ("yellow") line — used in tooltip copy that +// refers to that colour instead of naming it, so the swatch always matches the chart. +function WarningSwatch() { + return ( + + ); +} + +type QueueHeaderTile = { + id: string; + label: string; + /** Info-icon copy explaining what the chart shows, rendered next to the card title. */ + description: ReactNode; + color: string; + /** Optional inline legend rendered below the card title: a fixed set of {colored square, label} + * entries, for charts where a colour (e.g. the orange warning line) needs explaining. */ + legend?: Array<{ color: string; label: string }>; + query: string; + /** Formats a single bucket's value in the chart tooltip. */ + formatValue?: (value: number) => string; + /** Formats the y-axis tick labels. Without it the axis shows raw numbers (bad for durations + * in ms or percent scales). Passed through to Chart.Line's yAxisProps.tickFormatter. */ + formatAxis?: (value: number) => string; + /** Hover tooltip explaining the headline readout next to the title (e.g. what "9% of current + * period" means). Without it the readout has no tooltip. */ + totalTooltip?: string; + /** Turns one row per bucket into the per-bucket points the chart draws. A null value is a + * bucket the metric has nothing to say about, and the line breaks there rather than reading 0. */ + derive: (rows: MetricTileRow[]) => { + points: TilePoint[]; + total: number; + formatTotal?: (total: number) => string; + totalClassName?: string; }; - extraConcurrency: number; - extraUnallocatedConcurrency: number; - maxQuota: number; - canManageBilling: boolean; - disabled: boolean; -}) { - const showSelfServe = useShowSelfServe(); - const lastSubmission = useActionData(); - const [form, fields] = useForm({ - id: "purchase-concurrency", - // TODO: type this - lastResult: lastSubmission as any, - onValidate({ formData }) { - return parseWithZod(formData, { schema: FormSchema }); - }, - shouldRevalidate: "onSubmit", - }); - const { amount } = fields; + /** + * Optional second query, run at the range's natural bucket width, that owns the headline readout. + * For a headline that is not invariant to bucket width — a share of buckets, or a percentile, + * as opposed to a max over gauges — deriving it from the plotted rows would move the number + * whenever this tile's floor widens them. Only requested while the floor is actually widening + * anything; on ranges whose natural width is already at or above the floor the chart's own rows + * are identical and are reused. + */ + readout?: { + query: string; + derive: (rows: MetricTileRow[]) => { + total: number; + formatTotal?: (total: number) => string; + totalClassName?: string; + }; + }; +}; - const [amountValue, setAmountValue] = useState(extraConcurrency); - const navigation = useNavigation(); - const isLoading = navigation.state !== "idle" && navigation.formMethod === "POST"; +function tileNumber(value: number | string | null): number { + const n = typeof value === "number" ? value : Number(value); + return Number.isFinite(n) ? n : 0; +} - // Close the panel, when we've succeeded - // This is required because a redirect to the same path doesn't clear state - const [searchParams, setSearchParams] = useSearchParams(); - const purchaseSucceeded = Boolean(searchParams.get("success")); - const [open, setOpen] = useState(false); - useEffect(() => { - if (purchaseSucceeded) { - // oxlint-disable-next-line react/set-state-in-effect -- This effect intentionally synchronizes route state after an external or lifecycle change. - setOpen(false); - setSearchParams((s) => { - s.delete("success"); - return s; +function tileTimeToMs(value: number | string | null): number { + const s = String(value).replace(" ", "T"); + return Date.parse(s.endsWith("Z") ? s : `${s}Z`); +} + +/** Peak of a series, ignoring the buckets it has nothing to say about. */ +function peakOf(points: TilePoint[]): number { + return points.reduce((max, p) => (p.value === null ? max : Math.max(max, p.value)), 0); +} + +const SCHEDULING_DELAY_QUERY = `SELECT timeBucket() AS t,\n round(quantilesTDigestMerge(0.5, 0.9, 0.95, 0.99)(wait_quantiles)[3]) AS p95,\n sum(wait_ms_count) AS samples\nFROM env_metrics\nGROUP BY t\nORDER BY t`; + +const THROTTLED_QUERY = `SELECT timeBucket() AS t,\n sum(throttled_count) AS throttled\nFROM env_metrics\nGROUP BY t\nORDER BY t`; + +const QUEUE_HEADER_TILES: QueueHeaderTile[] = [ + { + id: "saturation", + label: "Env saturation", + description: ( + <> + How much of the environment's concurrency is in use. Turns above 100%, + when it's into burst capacity. + + ), + color: "var(--color-queues-chart)", + legend: [ + { color: "var(--color-queues-chart)", label: "Saturation" }, + { color: "var(--color-warning)", label: "Over limit" }, + ], + query: `SELECT timeBucket() AS t,\n max(max_env_running) AS running,\n max(max_env_limit) AS env_limit\nFROM env_metrics\nGROUP BY t\nORDER BY t`, + formatValue: (v) => (v > 100 ? `${v}% — over the environment limit` : `${v}%`), + formatAxis: (v) => `${v}%`, + derive: (rows) => { + const points = rows.map((r) => { + const limit = tileNumber(r.env_limit); + return { + bucket: tileTimeToMs(r.t), + value: limit > 0 ? Math.round((tileNumber(r.running) / limit) * 100) : 0, + }; }); - } - }, [purchaseSucceeded, setSearchParams]); + return { + points, + total: peakOf(points), + formatTotal: (v) => `${v}% peak`, + }; + }, + }, + { + id: "backlog", + label: "Backlog", + description: "How many runs are waiting across the environment, over time.", + color: "var(--color-queues-chart)", + query: `SELECT timeBucket() AS t,\n max(max_env_queued) AS queued\nFROM env_metrics\nGROUP BY t\nORDER BY t`, + derive: (rows) => { + const points = rows.map((r) => ({ + bucket: tileTimeToMs(r.t), + value: tileNumber(r.queued), + })); + return { + points, + total: peakOf(points), + formatTotal: (v) => `${v.toLocaleString()} peak`, + }; + }, + }, + { + id: "p95", + label: "Scheduling delay p95", + description: ( + <> + How long runs wait before they start (95% start faster than this). Turns {" "} + above 1 minute. + + ), + totalTooltip: "The worst p95 in the selected window.", + color: "var(--color-queues-chart)", + legend: [ + { color: "var(--color-queues-chart)", label: "p95" }, + { color: "var(--color-warning)", label: "Over 1 min" }, + ], + query: SCHEDULING_DELAY_QUERY, + formatValue: formatWaitMs, + formatAxis: formatWaitMs, + derive: (rows) => { + const points = rows.map((r) => ({ + bucket: tileTimeToMs(r.t), + value: tileNumber(r.samples) > 0 ? tileNumber(r.p95) : null, + })); + return { points, total: peakOf(points) }; + }, + readout: { + query: SCHEDULING_DELAY_QUERY, + /** + * Merging quantile states over a wider bucket yields a p95 between the sub-buckets' own, so + * the worst p95 has to be read at the range's natural width or a burst of slow starts shorter + * than the plotted bucket is averaged away. Unlike the gauges, whose max of maxes is the same + * at any width. + */ + derive: (rows) => { + const worst = rows.reduce( + (max, r) => (tileNumber(r.samples) > 0 ? Math.max(max, tileNumber(r.p95)) : max), + 0 + ); + return { + total: worst, + formatTotal: (v) => (v > 0 ? formatWaitMs(v) : "–"), + totalClassName: worst >= 60_000 ? "text-warning" : undefined, + }; + }, + }, + }, + { + id: "throttled", + label: "Throttled", + description: "How often runs were held back by a limit.", + totalTooltip: "The share of the selected window with at least one blocked dequeue.", + color: "var(--color-queues-chart)", + legend: [{ color: "var(--color-warning)", label: "Throttled" }], + query: THROTTLED_QUERY, + derive: (rows) => { + const points = rows.map((r) => ({ + bucket: tileTimeToMs(r.t), + value: tileNumber(r.throttled), + })); + return { points, total: peakOf(points) }; + }, + readout: { + query: THROTTLED_QUERY, + /** + * Share of the window that saw any throttling. A raw event sum isn't interpretable (it + * scales with poll rate and window length); the fraction of buckets with a throttle is. + * Gap fill zero-fills this counter, so every bucket in the window is present and the row + * count is the honest denominator. + */ + derive: (rows) => { + const nonzero = rows.filter((r) => tileNumber(r.throttled) > 0).length; + const pct = rows.length > 0 ? Math.round((nonzero / rows.length) * 100) : 0; + return { + total: pct, + formatTotal: (v) => `${v}% of current period`, + totalClassName: pct > 0 ? "text-warning" : undefined, + }; + }, + }, + }, +]; + +/** + * Bucket floor shared by every hero tile. Scheduling delay and throttling are event-driven, so at + * the 10-second width a short range would otherwise pick, most buckets hold no samples at all. One + * floor for all four keeps their x-axes identical, which the shared hover crosshair relies on. + */ +const HERO_CHART_MIN_BUCKET_SECONDS = 60; + +type TileTimeRange = MetricResourceTimeRange; + +// Full-size env metric chart rendered inside a ChartCard. Same data path as before +// (client-side TRQL via useMetricResourceQuery with fillGaps), drawn as a line that +// participates in the shared hover + drag-to-zoom of the enclosing ChartSyncProvider. +function QueueEnvMetricChart({ + tile, + timeRange, + referenceLines, + thresholdStroke, + warningOverlay, + solidWarning = false, +}: { + tile: QueueHeaderTile; + timeRange: TileTimeRange; + referenceLines?: Array<{ + y: number; + label?: string; + labelPlacement?: "inside" | "outside"; + }>; + thresholdStroke?: { value: number; aboveColor: string }; + warningOverlay?: { threshold: number }; + /** When set, the ENTIRE line turns warning-coloured if the series is ever non-zero (used for + * throttling: any throttle in the window colours the whole line). Mutually exclusive with the + * per-bucket warningOverlay. */ + solidWarning?: boolean; +}) { + const organization = useOrganization(); + const project = useProject(); + const environment = useEnvironment(); + + const sharedOptions = { + organizationId: organization.id, + projectId: project.id, + environmentId: environment.id, + timeRange, + defaultPeriod: QUEUE_METRICS_DEFAULT_PERIOD, + fillGaps: true, + }; - const state = updateState({ - value: amountValue, - existingValue: extraConcurrency, - quota: maxQuota, - extraUnallocatedConcurrency, + const { rows, showLoading, failed } = useMetricResourceQuery(tile.query, { + ...sharedOptions, + minBucketSeconds: HERO_CHART_MIN_BUCKET_SECONDS, }); - const changeClassName = - state === "decrease" ? "text-error" : state === "increase" ? "text-success" : undefined; - - const title = extraConcurrency === 0 ? "Purchase extra concurrency" : "Add/remove concurrency"; - - if (!showSelfServe) { - return ( - Request more} - /> - ); - } + + const derived = tile.derive(rows); + const points = derived.points; + + const plottedBucketMs = points.length > 1 ? points[1]!.bucket - points[0]!.bucket : 0; + const floorWidenedBuckets = + plottedBucketMs > 0 && plottedBucketMs <= HERO_CHART_MIN_BUCKET_SECONDS * 1000; + const readoutQuery = tile.readout && floorWidenedBuckets ? tile.readout.query : ""; + const readoutResult = useMetricResourceQuery(readoutQuery, sharedOptions); + + const { total, formatTotal, totalClassName } = tile.readout + ? tile.readout.derive(readoutQuery ? readoutResult.rows : rows) + : derived; + + // Same point shape the shared axis/tooltip helpers expect. + const data = points + .map((p) => ({ bucket: p.bucket, [tile.id]: p.value })) + .filter((p) => Number.isFinite(p.bucket)); + + // Whole-line warning colour when the series was ever non-zero (throttling: one throttle in the + // window colours the entire line). Otherwise the tile's normal colour. + const wholeLineWarning = solidWarning && total > 0; + const lineColor = wholeLineWarning ? "var(--color-warning)" : tile.color; + + const chartConfig = useMemo( + () => ({ [tile.id]: { label: tile.label, color: lineColor } }), + [tile.id, tile.label, lineColor] + ); + + const { tickFormatter, tooltipLabelFormatter } = buildActivityTimeAxis(data); + const hasData = data.length > 0 && data.some((p) => Number(p[tile.id] ?? 0) > 0); + + // Peak readout lives in the card title (ChartCard has no dedicated value slot). A zero/empty + // total renders no readout at all (skipping "0% peak", "0 peak", "0" and the p95 "–" placeholder) + // so the card title stands alone until there's a non-zero value to show. + const readoutLoading = tile.readout ? readoutResult.showLoading : showLoading; + const readoutFailed = tile.readout ? readoutResult.failed : failed; + const peak = readoutLoading ? ( + + ) : readoutFailed || total === 0 ? null : formatTotal ? ( + formatTotal(total) + ) : ( + total.toLocaleString() + ); return ( - - - - - - {title} -
-
- - You can purchase bundles of {concurrencyPricing.stepSize} concurrency for{" "} - {formatCurrency(concurrencyPricing.centsPerStep / 100, false)}/month. Or you can - remove any extra concurrency after you have unallocated it from your environments - first. - -
- - - setAmountValue(Number(e.target.value))} - disabled={isLoading} + + + + {tile.label} + + + {peak != null ? ( + tile.totalTooltip && !readoutLoading ? ( + + {peak} + + } + content={tile.totalTooltip} + className="max-w-[230px]" + disableHoverableContent /> - {amount.errors} - {form.errors} - -
- {state === "need_to_increase_unallocated" ? ( -
- - You need to unallocate{" "} - {formatNumber(extraConcurrency - amountValue - extraUnallocatedConcurrency)} more - concurrency from your environments in order to remove{" "} - {formatNumber(extraConcurrency - amountValue)} concurrency from your account. - -
- ) : state === "above_quota" ? ( -
- - Currently you can only have up to {maxQuota} extra concurrency. Send a request - below to lift your current limit. We'll get back to you soon. - -
- ) : ( -
-
- Summary - Total -
-
- - {formatNumber(extraConcurrency)}{" "} - current total - - - {formatCurrency( - (extraConcurrency * concurrencyPricing.centsPerStep) / - concurrencyPricing.stepSize / - 100, - true - )} - -
-
- - ({simplur`${extraConcurrency / concurrencyPricing.stepSize} bundle[|s]`}) - - /mth + ) : ( + + {peak} + + ) + ) : null} + + {tile.legend && (showLoading || hasData) ? ( + + {tile.legend.map((item) => ( + + + {item.label} + + ))} + + ) : null} + + } + > + {showLoading ? ( + + ) : failed ? ( +
+ Unable to load metrics +
+ ) : hasData ? ( + + + + ) : ( +
+ No activity +
+ )} + + ); +} + +function QueueMetricChartSkeleton() { + return ( +
+ {Array.from({ length: 42 }).map((_, i) => ( +
+ ))} +
+ ); +} + +/** Health as a stock Badge: color carries the state, w-fit keeps it content-width. */ +type QueueHealth = { + paused: boolean; + running: number; + queued: number; + limit: number; +}; + +type QueueHealthLabel = "Paused" | "At capacity" | "Backlogged" | "Active" | "Idle"; + +// Single source of truth for the queue health decision, shared by the badge and the table's +// health-column sort so the sorted order always matches the labels shown. +function queueHealthLabel({ paused, running, queued, limit }: QueueHealth): QueueHealthLabel { + if (paused) return "Paused"; + if (isQueueAtCapacity({ running, queued, limit })) return "At capacity"; + if (queued > 0) return "Backlogged"; + if (running > 0) return "Active"; + return "Idle"; +} + +// Tint + colored text, sized like the error status chips (see ErrorStatusBadge). +const QUEUE_HEALTH_STYLES: Record = { + Paused: "bg-warning/10 text-warning system:bg-warning system:text-white", + "At capacity": "bg-warning/10 text-warning system:bg-warning system:text-white", + Backlogged: "bg-blue-500/10 text-blue-500 system:bg-blue-500 system:text-white", + Active: "bg-success/10 text-success system:bg-success system:text-white", + Idle: "bg-charcoal-500/10 text-text-dimmed system:bg-charcoal-500 system:text-white", +}; + +function QueueHealthBadge(health: QueueHealth) { + const label = queueHealthLabel(health); + return ( + + {label} + + ); +} + +// The metrics-row key a queue is stored under (task queues are prefixed `task/`). +function queueMetricsKey(queue: { type: string; name: string }): string { + return `${queue.type === "task" ? "task/" : ""}${queue.name}`; +} + +function formatWaitMs(ms: number): string { + if (ms < 1000) return `${Math.round(ms)}ms`; + if (ms < 60_000) return `${(ms / 1000).toFixed(1)}s`; + if (ms < 3_600_000) return `${(ms / 60_000).toFixed(1)}m`; + return `${(ms / 3_600_000).toFixed(1)}h`; +} + +// Drop a trailing ".00" from whole percentages so "50.00" reads as "50" but "12.50" is preserved. +function formatOverridePercent(percent: number): string { + return Number.isInteger(percent) ? percent.toString() : percent.toFixed(2).replace(/\.?0+$/, ""); +} + +// Classic Queues page, restored verbatim from before the Queue Metrics feature. Rendered +// when queueMetricsUiEnabled is off so a gated org sees exactly the pre-metrics UI. +function ClassicQueuesView() { + const { environment, queues, pagination, hasFilters, autoReloadPollIntervalMs, canWriteTasks } = + useTypedLoaderData(); + + const organization = useOrganization(); + const project = useProject(); + const env = useEnvironment(); + const plan = useCurrentPlan(); + + useAutoRevalidate({ interval: autoReloadPollIntervalMs, onFocus: true }); + + const { limitStatus, limitClassName } = getEnvConcurrencyLimitStatus(environment); + + return ( + + + + + + + + Queues docs + + + + + +
+
+ paused : undefined} + animate + accessory={ +
+ {environment.runsEnabled && + env.pauseSource !== ENVIRONMENT_PAUSE_SOURCE_BILLING_LIMIT ? ( + + ) : null} +
-
- - {state === "increase" ? "+" : null} - {formatNumber(amountValue - extraConcurrency)} - - - {state === "increase" ? "+" : null} - {formatCurrency( - ((amountValue - extraConcurrency) * concurrencyPricing.centsPerStep) / - concurrencyPricing.stepSize / - 100, - true + } + valueClassName={env.paused ? "text-warning tabular-nums" : "tabular-nums"} + compactThreshold={1000000} + /> + -
-
- - ( - {simplur`${ - (amountValue - extraConcurrency) / concurrencyPricing.stepSize - } bundle[|s]`}{" "} - @ {formatCurrency(concurrencyPricing.centsPerStep / 100, true)}/mth) + > + Including {environment.running - environment.concurrencyLimit} burst runs{" "} + - /mth -
-
- - {formatNumber(amountValue)} new total - - - {formatCurrency( - (amountValue * concurrencyPricing.centsPerStep) / - concurrencyPricing.stepSize / - 100, - true + ) : limitStatus === "limit" ? ( + "At concurrency limit" + ) : undefined + } + accessory={ + + } + compactThreshold={1000000} + /> + 1 ? ( + -
-
- - ({simplur`${amountValue / concurrencyPricing.stepSize} bundle[|s]`}) + > + Burst limit {environment.burstFactor * environment.concurrencyLimit}{" "} + - /mth -
-
- )} + ) : undefined + } + accessory={ + plan ? ( + plan?.v3Subscription?.plan?.limits.concurrentRuns.canExceed ? ( + + Increase limit + + ) : ( + + Increase limit + + ) + ) : null + } + />
- - - - - ) : state === "decrease" || state === "need_to_increase_unallocated" ? ( - <> - -
+
+ + ); } -function updateState({ - value, - existingValue, - quota, - extraUnallocatedConcurrency, +function BurstFactorTooltip({ + environment, }: { - value: number; - existingValue: number; - quota: number; - extraUnallocatedConcurrency: number; -}): "no_change" | "increase" | "decrease" | "above_quota" | "need_to_increase_unallocated" { - if (value === existingValue) return "no_change"; - if (value < existingValue) { - const difference = existingValue - value; - if (difference > extraUnallocatedConcurrency) { - return "need_to_increase_unallocated"; - } - return "decrease"; - } - if (value > quota) return "above_quota"; - return "increase"; + environment: { burstFactor: number; concurrencyLimit: number }; +}) { + return ( + + ); } + +const limitTooltip = ( + <> + + How many runs can execute at once.{" "} + + + 1 per key (3) means 1 run per concurrency + key, but at most 3 runs across all keys.{" "} + 2 total caps every run of the queue, + keyed or not. Set using the concurrency{" "} + option in your code. + + +); diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.queues_.$queueParam/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.concurrency_.$queueParam/route.tsx similarity index 93% rename from apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.queues_.$queueParam/route.tsx rename to apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.concurrency_.$queueParam/route.tsx index 33e49bb6858..7acd4caf67d 100644 --- a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.queues_.$queueParam/route.tsx +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.concurrency_.$queueParam/route.tsx @@ -1,6 +1,7 @@ import { type LoaderFunctionArgs } from "@remix-run/server-runtime"; import { useCallback, useEffect, useMemo, useRef, useState, type ReactNode } from "react"; import type { QueueItem } from "@trigger.dev/core/v3/schemas"; +import type { QueueLimits } from "~/components/queues/queue-limits"; import { typedjson, useTypedLoaderData } from "remix-typedjson"; import { z } from "zod"; import { MainCenteredContainer, PageContainer } from "~/components/layout/AppLayout"; @@ -98,7 +99,7 @@ export const handle: Handle = { export const meta = pageMeta(({ data, params }) => [ data?.queue?.name ?? params.queueParam ?? "Queue", - "Queues", + "Concurrency", ]); const ParamsSchema = EnvironmentParamSchema.extend({ queueParam: z.string() }); @@ -129,7 +130,10 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { }); const canWriteTasks = auth.ok && auth.ability.can("write", { type: "tasks" }); - const retrieve = await new QueueRetrievePresenter().call({ environment, queueInput: queueParam }); + const retrieve = await new QueueRetrievePresenter().call({ + environment, + queueInput: queueParam, + }); if (!retrieve.success) { throw new Response(undefined, { status: 404, statusText: "Queue not found" }); } @@ -193,7 +197,7 @@ export const action = dashboardAction( const { organizationSlug, projectParam, envParam, queueParam } = ParamsSchema.parse(params); const url = new URL(request.url); - const redirectPath = `/orgs/${organizationSlug}/projects/${projectParam}/env/${envParam}/queues/${queueParam}${url.search}`; + const redirectPath = `/orgs/${organizationSlug}/projects/${projectParam}/env/${envParam}/concurrency/${queueParam}${url.search}`; if (request.method.toLowerCase() !== "post") { return redirectWithErrorMessage(redirectPath, request, "Wrong method"); @@ -289,7 +293,7 @@ export default function Page() { // The Concurrency keys tab exists only for queues with key activity: live keys in the // ckIndex, or nonzero CK history in the selected range (one cached scalar query decides). const { rows: gateRows, showLoading: gateLoading } = useQueueMetric( - `SELECT max(max_ck_backlogged) AS peak_keys, max(max_ck_wait_ms) AS peak_wait\nFROM queue_metrics`, + `SELECT max(max_ck_backlogged) AS peak_keys, max(max_ck_wait_ms) AS peak_wait\nFROM concurrency_metrics`, { ids, timeRange, queueName: fullName } ); const gateRow = gateRows[0]; @@ -308,14 +312,14 @@ export default function Page() { paused: queue.paused, running: queue.running, queued: queue.queued, - limit: queue.concurrencyLimit ?? environmentConcurrencyLimit, + limit: queue.limits.perKey.current ?? environmentConcurrencyLimit, oldestWaitMs, }); return ( - + {/* Paused-queue banner — mirrors the environment-paused banner (OrgBanner) at the top of the page when this individual queue is paused. */} @@ -416,7 +420,12 @@ export default function Page() { ) ) : ( - + )} @@ -460,10 +469,12 @@ function OverviewCharts({ ids, timeRange, queueName, + hasTotalLimit, }: { ids: Ids; timeRange: TimeRangeParams; queueName: string; + hasTotalLimit: boolean; }) { const zoomToTimeFilter = useZoomToTimeFilter(); return ( @@ -480,7 +491,7 @@ function OverviewCharts({ } showLegend className="aspect-[2/1]" - query={`SELECT timeBucket() AS t, max(max_running) AS running, max(max_limit) AS limit\nFROM queue_metrics\nGROUP BY t\nORDER BY t`} + query={`SELECT timeBucket() AS t, max(max_running) AS running, max(max_limit) AS limit\nFROM concurrency_metrics\nGROUP BY t\nORDER BY t`} fillGaps minBucketSeconds={SYNCED_CHART_MIN_BUCKET_SECONDS} ids={ids} @@ -503,11 +514,43 @@ function OverviewCharts({ // leading zeros so the reference line doesn't start with a false 0→limit step. carryBackfill={["limit"]} /> + {hasTotalLimit ? ( + + Runs in flight across ALL concurrency keys ( + ) versus the queue's total limit ( + + ). + + } + showLegend + className="aspect-[2/1]" + query={`SELECT timeBucket() AS t, max(max_total_running) AS running, least(nullIf(max(max_total_limit), 0), max(max_env_limit)) AS cap, max(max_env_limit) AS sampled\nFROM concurrency_metrics\nGROUP BY t\nORDER BY t`} + fillGaps + minBucketSeconds={SYNCED_CHART_MIN_BUCKET_SECONDS} + ids={ids} + timeRange={timeRange} + queueName={queueName} + series={[ + { key: "cap", label: "Total limit", color: COLORS.limit }, + { key: "running", label: "Running", color: COLORS.running }, + ]} + thresholdStroke={{ + series: "running", + valueFromSeries: "cap", + aboveColor: "var(--color-warning)", + }} + carryBackfill={["cap"]} + carryBackfillGuard="sampled" + /> + ) : null} } className="aspect-[2/1] sm:col-span-2 sm:aspect-[4/1]" - query={`SELECT timeBucket() AS t, sum(throttled_count) AS throttled\nFROM queue_metrics\nGROUP BY t\nORDER BY t`} + query={`SELECT timeBucket() AS t, sum(throttled_count) AS throttled\nFROM concurrency_metrics\nGROUP BY t\nORDER BY t`} fillGaps minBucketSeconds={SYNCED_CHART_MIN_BUCKET_SECONDS} ids={ids} @@ -684,7 +727,7 @@ function ConcurrencyKeyCharts({ title="Keys with backlog" info="Keys with runs waiting at once." className="aspect-[2/1]" - query={`SELECT timeBucket() AS t, max(max_ck_backlogged) AS keys\nFROM queue_metrics\nGROUP BY t\nORDER BY t`} + query={`SELECT timeBucket() AS t, max(max_ck_backlogged) AS keys\nFROM concurrency_metrics\nGROUP BY t\nORDER BY t`} fillGaps ids={ids} timeRange={timeRange} @@ -706,7 +749,7 @@ function ConcurrencyKeyCharts({ ) : null } className="aspect-[2/1]" - query={`SELECT timeBucket() AS t, max(max_ck_wait_ms) AS wait\nFROM queue_metrics\nGROUP BY t\nORDER BY t`} + query={`SELECT timeBucket() AS t, max(max_ck_wait_ms) AS wait\nFROM concurrency_metrics\nGROUP BY t\nORDER BY t`} fillGaps ids={ids} timeRange={timeRange} @@ -768,7 +811,7 @@ type GroupedKeyChartProps = { // search can match keys outside the top 8; then filter by the search and keep the top 8 of those. function GroupedKeyChartCard(props: GroupedKeyChartProps) { const { rows, showLoading, failed } = useQueueMetric( - `SELECT concurrency_key, ${props.rankExpr} AS peak\nFROM queue_metrics_by_key\nGROUP BY concurrency_key\nORDER BY peak DESC\nLIMIT 50`, + `SELECT concurrency_key, ${props.rankExpr} AS peak\nFROM concurrency_metrics_by_key\nGROUP BY concurrency_key\nORDER BY peak DESC\nLIMIT 50`, { ids: props.ids, timeRange: props.timeRange, queueName: props.queueName } ); const keyFilter = props.keyFilter; @@ -796,7 +839,7 @@ function GroupedKeySeries({ }: GroupedKeyChartProps & { keys: string[] }) { const inList = keys.map((k) => `'${trqlString(k)}'`).join(", "); const { rows, showLoading, failed } = useQueueMetric( - `SELECT timeBucket() AS t, concurrency_key, ${seriesExpr} AS v\nFROM queue_metrics_by_key\nWHERE concurrency_key IN (${inList})\nGROUP BY t, concurrency_key\nORDER BY t`, + `SELECT timeBucket() AS t, concurrency_key, ${seriesExpr} AS v\nFROM concurrency_metrics_by_key\nWHERE concurrency_key IN (${inList})\nGROUP BY t, concurrency_key\nORDER BY t`, { ids, timeRange, queueName, fillGaps } ); @@ -933,6 +976,7 @@ function useConcurrencyKeys(opts: { interval: 30_000, onLoad: false, onFocus: true, + pauseWhenHidden: true, callback: load, }); @@ -1059,7 +1103,7 @@ function KeyDrilldown({ } className="aspect-[2/1]" - query={`SELECT timeBucket() AS t, max(max_queued) AS queued, max(max_running) AS running\nFROM queue_metrics_by_key\nWHERE ${pin}\nGROUP BY t\nORDER BY t`} + query={`SELECT timeBucket() AS t, max(max_queued) AS queued, max(max_running) AS running\nFROM concurrency_metrics_by_key\nWHERE ${pin}\nGROUP BY t\nORDER BY t`} fillGaps minBucketSeconds={SYNCED_CHART_MIN_BUCKET_SECONDS} ids={ids} @@ -1074,7 +1118,7 @@ function KeyDrilldown({ 0, round(sum(wait_ms_sum) / sum(wait_ms_count)), 0) AS wait, sum(wait_ms_count) AS samples\nFROM queue_metrics_by_key\nWHERE ${pin}\nGROUP BY t\nORDER BY t`} + query={`SELECT timeBucket() AS t, if(sum(wait_ms_count) > 0, round(sum(wait_ms_sum) / sum(wait_ms_count)), 0) AS wait, sum(wait_ms_count) AS samples\nFROM concurrency_metrics_by_key\nWHERE ${pin}\nGROUP BY t\nORDER BY t`} fillGaps minBucketSeconds={SYNCED_CHART_MIN_BUCKET_SECONDS} sampleCountColumn="samples" @@ -1122,7 +1166,7 @@ function QueueStats({ }: { // Carries the percent override source-of-truth (not part of the shared QueueItem contract) so the // override dialog reopens in percent mode for percent-based overrides. - queue: QueueItem & { concurrencyLimitOverridePercent: number | null }; + queue: QueueItem & { concurrencyLimitOverridePercent: number | null; limits: QueueLimits }; environmentConcurrencyLimit: number; queuedRunsPath: string; oldestWaitMs: number | null; @@ -1130,18 +1174,21 @@ function QueueStats({ timeRange: TimeRangeParams; queueName: string; }) { - const { rows } = useQueueMetric(`SELECT max(max_queued) AS peak_queued\nFROM queue_metrics`, { - ids, - timeRange, - queueName, - }); + const { rows } = useQueueMetric( + `SELECT max(max_queued) AS peak_queued\nFROM concurrency_metrics`, + { + ids, + timeRange, + queueName, + } + ); const peakQueued = rows[0] ? toNumber(rows[0].peak_queued) : 0; // Latest gauges from ClickHouse, polled every 15s so the live blocks keep ticking after first // paint. Read the newest bucket (largest t); until the first poll lands liveRows is empty and the // *Live values stay null, so the blocks show the loader values instead of flashing 0. const { rows: liveRows, responseReceivedAt } = useQueueMetric( - `SELECT timeBucket() AS t, max(max_running) AS running, max(max_queued) AS queued, max(max_limit) AS q_limit, max(max_ck_wait_ms) AS ck_wait FROM queue_metrics GROUP BY t ORDER BY t`, + `SELECT timeBucket() AS t, max(max_running) AS running, max(max_queued) AS queued, max(max_limit) AS q_limit, max(max_ck_wait_ms) AS ck_wait FROM concurrency_metrics GROUP BY t ORDER BY t`, { ids, timeRange: { period: "15m", from: null, to: null }, @@ -1172,7 +1219,7 @@ function QueueStats({ const queuedDisplay = queuedLive ?? queue.queued; // Limit is queue config, not a live signal: keep the loader's value. Only if the loader had none // do we fall back to the CH gauge for display. - const limitDisplay = queue.concurrencyLimit ?? (limitLive || null); + const limitDisplay = queue.limits.perKey.current ?? (limitLive || null); // Keyed queues report head-of-line wait via CH (max_ck_wait_ms); use it as the live headline when // present. Non-keyed queues have no CH signal, so they stay on the loader value. const oldestWaitDisplayMs = ckWaitLive !== null && ckWaitLive > 0 ? ckWaitLive : oldestWaitMs; diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.limits/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.limits/route.tsx index e8fb74da399..045c67e592b 100644 --- a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.limits/route.tsx +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.limits/route.tsx @@ -43,7 +43,7 @@ import { requireUserId } from "~/services/session.server"; import { cn } from "~/utils/cn"; import { formatNumber } from "~/utils/numberFormatter"; import { - concurrencyPath, + concurrencyLimitsPath, docsPath, EnvironmentParamSchema, organizationBillingPath, @@ -153,7 +153,7 @@ export default function Page() { {/* Concurrency Section */} {/* Rate Limits Section */} @@ -231,7 +231,7 @@ function CurrentPlanSection({ ); } -function ConcurrencySection({ concurrencyPath }: { concurrencyPath: string }) { +function ConcurrencySection({ concurrencyLimitsPath }: { concurrencyLimitsPath: string }) { return (
@@ -247,7 +247,7 @@ function ConcurrencySection({ concurrencyPath }: { concurrencyPath: string }) { Concurrency - + Manage concurrency diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.queues.$.ts b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.queues.$.ts new file mode 100644 index 00000000000..cc950eb6786 --- /dev/null +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.queues.$.ts @@ -0,0 +1,10 @@ +import { redirect, type LoaderFunctionArgs } from "@remix-run/server-runtime"; + +/** The queues page became the Concurrency page; old URLs (bookmarks, agent deep + * links, the queue detail path) redirect with their sub-path and query intact. */ +export const loader = async ({ params, request }: LoaderFunctionArgs) => { + const url = new URL(request.url); + const subPath = params["*"] ? `/${params["*"]}` : ""; + const base = url.pathname.replace(/\/queues(\/.*)?$/, "/concurrency"); + return redirect(`${base}${subPath}${url.search}`, 301); +}; diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.queues.ts b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.queues.ts new file mode 100644 index 00000000000..69646e8c781 --- /dev/null +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.queues.ts @@ -0,0 +1,8 @@ +import { redirect, type LoaderFunctionArgs } from "@remix-run/server-runtime"; + +/** The queues page became the Concurrency page; the bare old URL redirects with + * its query intact. */ +export const loader = async ({ request }: LoaderFunctionArgs) => { + const url = new URL(request.url); + return redirect(`${url.pathname.replace(/\/queues$/, "/concurrency")}${url.search}`, 301); +}; diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.queues/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.queues/route.tsx deleted file mode 100644 index 289a7bfe927..00000000000 --- a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.queues/route.tsx +++ /dev/null @@ -1,2080 +0,0 @@ -import { - ArrowUpCircleIcon, - BookOpenIcon, - ExclamationTriangleIcon, - PauseIcon, - PlayIcon, - RectangleStackIcon, -} from "@heroicons/react/20/solid"; -import { DialogClose } from "@radix-ui/react-dialog"; -import { Form, useNavigation } from "@remix-run/react"; -import { type LoaderFunctionArgs } from "@remix-run/server-runtime"; -import type { RuntimeEnvironmentType } from "@trigger.dev/database"; -import { useEffect, useMemo, useState, type ReactNode } from "react"; -import { QueuesIcon } from "~/assets/icons/QueuesIcon"; -import { typedjson, useTypedLoaderData } from "remix-typedjson"; -import { z } from "zod"; -import { ConcurrencyIcon } from "~/assets/icons/ConcurrencyIcon"; -import { RunsIcon } from "~/assets/icons/RunsIcon"; -import { AdminDebugTooltip } from "~/components/admin/debugTooltip"; -import { environmentFullTitle } from "~/components/environments/EnvironmentLabel"; -import { PageBody, PageContainer } from "~/components/layout/AppLayout"; -import { MetricsLayout } from "~/components/layout/MetricsLayout"; -import { Badge } from "~/components/primitives/Badge"; -import { Button, LinkButton } from "~/components/primitives/Buttons"; -import { Dialog, DialogContent, DialogHeader, DialogTrigger } from "~/components/primitives/Dialog"; -import { FormButtons } from "~/components/primitives/FormButtons"; -import { Header3 } from "~/components/primitives/Headers"; -import { NavBar, PageAccessories, PageTitle } from "~/components/primitives/PageHeader"; -import { PaginationControls } from "~/components/primitives/Pagination"; -import { Paragraph } from "~/components/primitives/Paragraph"; -import { PopoverMenuItem } from "~/components/primitives/Popover"; -import { SearchInput } from "~/components/primitives/SearchInput"; -import { Spinner } from "~/components/primitives/Spinner"; -import { - Table, - TableBody, - TableCell, - TableCellMenu, - TableHeader, - TableHeaderCell, - TableRow, -} from "~/components/primitives/Table"; -import { - InfoIconTooltip, - SimpleTooltip, - Tooltip, - TooltipContent, - TooltipProvider, - TooltipTrigger, -} from "~/components/primitives/Tooltip"; -import { TasksIcon } from "~/assets/icons/TasksIcon"; -import { QueueName } from "~/components/runs/v3/QueueName"; -import { env } from "~/env.server"; -import { useAutoRevalidate } from "~/hooks/useAutoRevalidate"; -import { useEnvironment } from "~/hooks/useEnvironment"; -import { useOrganization } from "~/hooks/useOrganizations"; -import { useProject } from "~/hooks/useProject"; -import { redirectWithErrorMessage, redirectWithSuccessMessage } from "~/models/message.server"; -import { findProjectBySlug } from "~/models/project.server"; -import { findEnvironmentBySlug } from "~/models/runtimeEnvironment.server"; -import { EnvironmentQueuePresenter } from "~/presenters/v3/EnvironmentQueuePresenter.server"; -import { QueueListPresenter } from "~/presenters/v3/QueueListPresenter.server"; -import { - QueueMetricsPresenter, - type QueueListMetric, -} from "~/presenters/v3/QueueMetricsPresenter.server"; -import { TimeFilter, timeFilterFromTo } from "~/components/runs/v3/SharedFilters"; -import { useSearchParams } from "~/hooks/useSearchParam"; -import { parseFiniteInt } from "~/utils/searchParams"; -import { MiniLineChart } from "~/components/metrics/MiniLineChart"; -import { buildActivityTimeAxis } from "~/components/primitives/charts/activityTimeAxis"; -import { Chart, type ChartConfig } from "~/components/primitives/charts/ChartCompound"; -import { ChartCard } from "~/components/primitives/charts/ChartCard"; -import { ChartSyncProvider } from "~/components/primitives/charts/ChartSyncContext"; -import { useZoomToTimeFilter } from "~/hooks/useZoomToTimeFilter"; -import { - useIsMetricResponseFresh, - useMetricResourceQuery, - type MetricResourceTimeRange, -} from "~/hooks/useMetricResourceQuery"; -import { logger } from "~/services/logger.server"; -import { rbac } from "~/services/rbac.server"; -import { resolveProjectAuthScope } from "~/services/projectAuthScope.server"; -import { dashboardAction } from "~/services/routeBuilders/dashboardBuilder"; -import { requireUserId } from "~/services/session.server"; -import { cn } from "~/utils/cn"; -import { ENVIRONMENT_PAUSE_SOURCE_BILLING_LIMIT } from "~/utils/environmentPauseSource"; -import { - concurrencyPath, - docsPath, - EnvironmentParamSchema, - v3BillingPath, - v3QueuePath, - v3RunsPath, -} from "~/utils/pathBuilder"; -import type { Handle } from "~/utils/handle"; -import { queuesAgentPageContext } from "~/components/dashboard-agent/suggested-prompts"; -import { WhenAgentUnavailable } from "~/components/dashboard-agent/WhenAgentUnavailable"; -import { PauseEnvironmentService } from "~/v3/services/pauseEnvironment.server"; -import { handleQueueMutationAction } from "~/models/queueMutation.server"; -import { - QueueOverrideConcurrencyButton, - QueuePauseResumeButton, -} from "~/components/queues/QueueControls"; -import { useCurrentPlan } from "../_app.orgs.$organizationSlug/route"; -import { BigNumber } from "~/components/metrics/BigNumber"; -import { canAccessQueueMetricsUi } from "~/v3/canAccessQueueMetricsUi.server"; -import { QueueAllocationPresenter } from "~/presenters/v3/QueueAllocationPresenter.server"; -import { - QUEUE_METRICS_DEFAULT_PERIOD, - QUEUE_METRICS_RETENTION_DAYS, - clampQueueMetricsPeriod, - clipQueueMetricsWindow, - queueMetricsPeriodFromRequest, - resolveQueueMetricsPeriod, - useRememberQueueMetricsPeriod, -} from "~/components/queues/queueMetricsPeriod"; -import { queueMetricsMaxPeriodDays } from "~/components/queues/queueMetricsPeriod.server"; -import { isQueueAtCapacity } from "~/components/queues/queue-thresholds"; -import { pageMeta } from "~/utils/pageTitle"; - -const SearchParamsSchema = z.object({ - query: z.string().optional(), - page: z.coerce.number().min(1).default(1), - period: z.string().optional(), - from: z.string().optional(), - to: z.string().optional(), - sort: z.enum(["busiest", "queued", "name"]).optional(), -}); - -// The live "Queued" / "Running" header blocks poll ClickHouse on a short cadence so they stay -// current after first paint. They read the env-wide gauges from env_metrics (the env-level rollup -// of queue_metrics, cheapest for a dimension-free query), always over a fixed 15m window regardless -// of the chart/table period, and are NOT scoped to the visible queue set (the blocks are env-wide). -const QUEUE_LIVE_BLOCKS_PERIOD = "15m"; -const QUEUE_LIVE_BLOCKS_QUERY = - "SELECT timeBucket() AS t, max(max_env_queued) AS env_queued, max(max_env_running) AS env_running FROM env_metrics GROUP BY t ORDER BY t"; -// Trust the ClickHouse gauge only while its newest bucket is this recent; otherwise fall back to -// the loader's Redis-exact live values (matches LIVE_GAUGE_FRESH_MS on the queue detail page / run -// inspector). -const LIVE_GAUGE_FRESH_MS = 90_000; - -export const handle: Handle = { - agentPageContext: (data) => queuesAgentPageContext(data), -}; - -export const meta = pageMeta("Queues"); - -export const loader = async ({ request, params }: LoaderFunctionArgs) => { - const userId = await requireUserId(request); - const { organizationSlug, projectParam, envParam } = EnvironmentParamSchema.parse(params); - - const url = new URL(request.url); - const { page, query, period, from, to, sort } = SearchParamsSchema.parse( - Object.fromEntries(url.searchParams) - ); - - const project = await findProjectBySlug(organizationSlug, projectParam, userId); - if (!project) { - throw new Response(undefined, { - status: 404, - statusText: "Project not found", - }); - } - - const environment = await findEnvironmentBySlug(project.id, envParam, userId); - if (!environment) { - throw new Response(undefined, { - status: 404, - statusText: "Environment not found", - }); - } - - const auth = await rbac.authenticateSession(request, { - userId, - organizationId: project.organizationId, - projectId: project.id, - }); - const canWriteTasks = auth.ok && auth.ability.can("write", { type: "tasks" }); - - // Per-org gate for the metrics UI. When off, this org gets the classic Queues page and - // no metrics query fires. - const queueMetricsUiEnabled = await canAccessQueueMetricsUi({ - request, - userId, - organizationSlug, - }); - - const maxPeriodDays = queueMetricsUiEnabled - ? await queueMetricsMaxPeriodDays(environment.organizationId) - : QUEUE_METRICS_RETENTION_DAYS; - const defaultPeriod = clampQueueMetricsPeriod( - queueMetricsPeriodFromRequest(request), - maxPeriodDays - ); - - try { - const queueListPresenter = new QueueListPresenter(); - const queues = await queueListPresenter.call({ - environment, - query, - page, - // Relevance ordering rides the metrics pipeline, so it is part of the gated UI. - sort: queueMetricsUiEnabled ? (sort ?? "busiest") : "name", - }); - - const environmentQueuePresenter = new EnvironmentQueuePresenter(); - - const autoReloadPollIntervalMs = env.QUEUES_AUTORELOAD_POLL_INTERVAL_MS; - - // Per-queue list metrics (Delay p95 + backlog sparkline columns) are SSR'd with the table. - // The environment header tiles are fetched client-side per card (see QueueEnvMetricChart) so a - // slow ClickHouse query never blocks the queues list from rendering. - let metrics: { - bucketStartMs: number; - bucketIntervalMs: number; - byQueue: Record; - } | null = null; - - if (queueMetricsUiEnabled) { - // Metrics are additive observability; a ClickHouse hiccup must not take down queue - // management. Fail open to metrics: null instead of bubbling to the page-level 400. - try { - const presenter = new QueueMetricsPresenter(); - const queueNames = queues.queues.map((q) => - q.type === "task" ? `task/${q.name}` : q.name - ); - const timeRange = clipQueueMetricsWindow( - timeFilterFromTo({ - period: - resolveQueueMetricsPeriod({ period, from, to, defaultPeriod, maxPeriodDays }) ?? - undefined, - from: parseFiniteInt(from), - to: parseFiniteInt(to), - defaultPeriod, - }), - maxPeriodDays - ); - const queueMetrics = - queueNames.length > 0 - ? await presenter.getQueueListMetrics({ - environment, - queueNames, - from: timeRange.from, - to: timeRange.to, - }) - : null; - if (queueMetrics) { - metrics = { - bucketStartMs: queueMetrics.bucketStartMs, - bucketIntervalMs: queueMetrics.bucketIntervalMs, - byQueue: Object.fromEntries(queueMetrics.byQueue), - }; - } - } catch (error) { - logger.warn("Queue list metrics unavailable, rendering without them", { - error, - organizationSlug: params.organizationSlug, - }); - } - } - - // Allocation summary (Environment limit + Allocated tiles) is additive; a presenter - // failure must not 400 the page, so fail open to null like the metrics block above. - let allocation: Awaited> | null = null; - if (queueMetricsUiEnabled) { - try { - allocation = await new QueueAllocationPresenter().call({ environment }); - } catch (error) { - logger.warn("Queue allocation summary unavailable, rendering without it", { - error, - organizationSlug: params.organizationSlug, - }); - } - } - - return typedjson({ - ...queues, - environment: await environmentQueuePresenter.call(environment), - autoReloadPollIntervalMs, - metrics, - allocation, - queueMetricsUiEnabled, - defaultPeriod, - maxPeriodDays, - canWriteTasks, - }); - } catch (error) { - console.error(error); - throw new Response(undefined, { - status: 400, - statusText: "Something went wrong, if this problem persists please contact support.", - }); - } -}; - -export const action = dashboardAction( - { - params: EnvironmentParamSchema, - context: (params) => resolveProjectAuthScope(params.organizationSlug, params.projectParam), - authorization: { - action: "write", - resource: { type: "tasks" }, - message: "With your current role, you can't manage queues.", - }, - }, - async ({ request, params, user }) => { - const userId = user.id; - if (request.method.toLowerCase() !== "post") { - return redirectWithErrorMessage( - `/orgs/${params.organizationSlug}/projects/${params.projectParam}/env/${params.envParam}/queues`, - request, - "Wrong method" - ); - } - - const { organizationSlug, projectParam, envParam } = EnvironmentParamSchema.parse(params); - - const project = await findProjectBySlug(organizationSlug, projectParam, userId); - if (!project) { - throw new Response(undefined, { - status: 404, - statusText: "Project not found", - }); - } - - const environment = await findEnvironmentBySlug(project.id, envParam, userId); - if (!environment) { - throw new Response(undefined, { - status: 404, - statusText: "Environment not found", - }); - } - - const formData = await request.formData(); - const action = formData.get("action"); - - const url = new URL(request.url); - const redirectPath = `/orgs/${organizationSlug}/projects/${projectParam}/env/${envParam}/queues${url.search}`; - - if (environment.archivedAt) { - return redirectWithErrorMessage(redirectPath, request, "This branch is archived"); - } - - // Per-queue actions (pause/resume/override/remove-override) are shared with the queue detail - // route, so they live in a helper that both routes call. - const queueMutation = await handleQueueMutationAction({ - request, - environment, - userId, - formData, - redirectPath, - }); - if (queueMutation) { - return queueMutation; - } - - switch (action) { - case "environment-pause": { - const pauseService = new PauseEnvironmentService(); - const result = await pauseService.call(environment, "paused"); - if (!result.success) { - return redirectWithErrorMessage(redirectPath, request, result.error); - } - return redirectWithSuccessMessage(redirectPath, request, "Environment paused"); - } - case "environment-resume": { - const resumeService = new PauseEnvironmentService(); - const result = await resumeService.call(environment, "resumed"); - if (!result.success) { - return redirectWithErrorMessage(redirectPath, request, result.error); - } - return redirectWithSuccessMessage(redirectPath, request, "Environment resumed"); - } - default: - return redirectWithErrorMessage(redirectPath, request, "Something went wrong"); - } - } -); - -// Derives the environment concurrency status ("limit" | "burst" | "within") and the matching -// text color from the current running count vs. the env limit and burst factor. Shared by both -// the classic and metrics views so the "Running" tile styling stays in sync. -function getEnvConcurrencyLimitStatus(environment: { - running: number; - concurrencyLimit: number; - burstFactor: number; -}) { - const limitStatus = - environment.running === environment.concurrencyLimit * environment.burstFactor - ? "limit" - : environment.running > environment.concurrencyLimit - ? "burst" - : "within"; - - const limitClassName = - limitStatus === "burst" ? "text-warning" : limitStatus === "limit" ? "text-error" : undefined; - - return { limitStatus, limitClassName }; -} - -export default function Page() { - // Per-org flag decides which whole page renders. Off => the classic Queues page, - // byte-for-byte the pre-metrics UI. Each branch is its own component (own hooks). - const { queueMetricsUiEnabled } = useTypedLoaderData(); - return queueMetricsUiEnabled ? : ; -} - -function QueuesWithMetricsView() { - const { - environment, - queues, - pagination, - totalQueues, - hasFilters, - autoReloadPollIntervalMs, - metrics, - allocation, - defaultPeriod, - maxPeriodDays, - canWriteTasks, - } = useTypedLoaderData(); - - const metricsByQueue = metrics?.byQueue ?? {}; - - const organization = useOrganization(); - const project = useProject(); - const env = useEnvironment(); - const plan = useCurrentPlan(); - - // The header tiles fetch client-side with the same period/from/to the TimeFilter writes. - const { value } = useSearchParams(); - const timeRange = { - period: resolveQueueMetricsPeriod({ - period: value("period"), - from: value("from"), - to: value("to"), - defaultPeriod, - maxPeriodDays, - }), - from: value("from") ?? null, - to: value("to") ?? null, - }; - useRememberQueueMetricsPeriod(value("period")); - - useAutoRevalidate({ interval: autoReloadPollIntervalMs, onFocus: true }); - - // Drag-to-zoom on either chart narrows the page's from/to search params, which the - // TimeFilter and the client-side metric queries both read (same wiring as the Agent page). - const zoomToTimeFilter = useZoomToTimeFilter(); - - // Live env-wide Queued/Running blocks. First paint uses the loader's Redis-exact values; from the - // first poll on we prefer ClickHouse so the blocks stay current without a full page revalidate. - // Empty rows (quiet env, or the very first fetch still in flight) fall back to the loader values, - // so we never flash a stale 0. Fixed 15m window, env-wide (no queue filter), CH-only recurring - // load; pauses while the tab is hidden (handled inside the hook). - const { rows: liveBlockRows, responseReceivedAt } = useMetricResourceQuery( - QUEUE_LIVE_BLOCKS_QUERY, - { - organizationId: organization.id, - projectId: project.id, - environmentId: env.id, - timeRange: { period: QUEUE_LIVE_BLOCKS_PERIOD, from: null, to: null }, - defaultPeriod: QUEUE_LIVE_BLOCKS_PERIOD, - fillGaps: false, - refreshIntervalMs: 15_000, - } - ); - const lastLiveBlockRow = - liveBlockRows.length > 0 ? liveBlockRows[liveBlockRows.length - 1] : null; - // Only trust the gauge while its newest bucket is fresh. A row painted from the hook's cache on - // client-side nav-back (responseCache), or a quiet env whose latest bucket is minutes old, must - // not override the loader's Redis-exact live values with a stale count. - const lastLiveBucketMs = lastLiveBlockRow ? tileTimeToMs(lastLiveBlockRow.t) : NaN; - const liveBlockIsFresh = useIsMetricResponseFresh( - responseReceivedAt, - lastLiveBucketMs, - LIVE_GAUGE_FRESH_MS - ); - const freshLiveBlockRow = lastLiveBlockRow && liveBlockIsFresh ? lastLiveBlockRow : null; - const envQueuedLive = freshLiveBlockRow - ? tileNumber(freshLiveBlockRow.env_queued) - : environment.queued; - const envRunningLive = freshLiveBlockRow - ? tileNumber(freshLiveBlockRow.env_running) - : environment.running; - - // Allocation summary tiles. The presenter computes the env-wide allocated total (sum of - // each queue's explicit limit clamped to the env limit) in a single aggregate query. - const envLimit = environment.concurrencyLimit; - const burstLimit = Math.round(envLimit * environment.burstFactor); - const allocated = allocation?.allocated ?? 0; - const allocationPct = envLimit > 0 ? Math.round((allocated / envLimit) * 100) : 0; - - // Running-block tinting (burst/limit) tracks the live running value, not the loader snapshot. - const { limitStatus, limitClassName } = getEnvConcurrencyLimitStatus({ - running: envRunningLive, - concurrencyLimit: environment.concurrencyLimit, - burstFactor: environment.burstFactor, - }); - - // Client-side, header-click sorting over the current page's rows. Server pagination and the - // default busiest order are unchanged; clearing a sort returns to that server order. - const queueRows = queues ?? []; - - return ( - - - - - - - - Queues docs - - - - - - {/* Filters — pinned bar directly under the NavBar. This row is page-wide only: Period is - the one control that changes the tiles and charts below, so it leads the row. Search - and pagination scope the table alone and live in that table's own bar instead. */} - -
- -
-
- {environment.runsEnabled && - env.pauseSource !== ENVIRONMENT_PAUSE_SOURCE_BILLING_LIMIT ? ( - - ) : null} -
- - - {/* Queued + Running + Allocated + Environment limit summary. Four stat tiles: the grid - derives its columns from the tile count (two-up, four-up from lg). The allocation - presenter fails open to null (a ClickHouse/PG hiccup mustn't take down the tiles), so - only the Allocated tile depends on it — the other three + controls always render, and - Allocated shows a "–" placeholder to keep the 4-tile grid shape stable. */} - - paused : undefined} - animate - accessory={ - - - - } - valueClassName={env.paused ? "text-warning tabular-nums" : "tabular-nums"} - compactThreshold={1000000} - /> - - Including {envRunningLive - environment.concurrencyLimit} burst runs{" "} - - - ) : limitStatus === "limit" ? ( - "At concurrency limit" - ) : undefined - } - accessory={ - - - - } - compactThreshold={1000000} - /> - - Allocated - {allocation ? ( - - ) : null} - - } - value={allocation ? allocated : undefined} - formattedValue={allocation ? undefined : "–"} - suffix={allocation ? `${allocationPct}% of the environment limit` : undefined} - suffixClassName="text-text-dimmed" - /> - 1 ? `bursts up to ${burstLimit}` : undefined} - suffixClassName="text-text-dimmed" - accessory={ - plan ? ( - plan?.v3Subscription?.plan?.limits.concurrentRuns.canExceed ? ( - - Increase limit - - ) : ( - - Increase limit - - ) - ) : undefined - } - /> - - - {hasFilters || totalQueues !== 0 ? ( - - - {QUEUE_HEADER_TILES.map((tile) => ( - 1 - ? [ - { - y: Math.round(environment.burstFactor * 100), - label: `Burst ${Math.round( - environment.concurrencyLimit * environment.burstFactor - )}`, - labelPlacement: "outside" as const, - }, - ] - : []), - ] - : undefined - } - // Saturation recolours the line above its 100% limit with a gradient split, so - // only the portion over the line is orange (the offset is derived from the line's - // own value range, so the split lands exactly at 100% regardless of domain - // padding). p95 and throttled use a per-bucket overlay: it retraces only the - // over-threshold stretches, so under-threshold buckets stay blue. - thresholdStroke={ - tile.id === "saturation" - ? { value: 100, aboveColor: "var(--color-warning)" } - : undefined - } - warningOverlay={ - tile.id === "p95" - ? { threshold: 60_000 } - : tile.id === "throttled" - ? // Integer counts: threshold 0 warns once a bucket has ≥1 throttle. - { threshold: 0 } - : undefined - } - /> - ))} - - - ) : null} - - - - - - } - > - {/* Default overflow-x-auto container so wide tables still scroll horizontally on - narrow viewports; the page (not this region) owns vertical scrolling. */} - - - - Name - Queued - Running - Limit - -

- Environment: uses the environment - limit of {environment.concurrencyLimit}. -

-

- User: a limit you set in your - code. -

-

- Override: a limit you set here or - via the API. -

- - } - > - Limited by - - Health - - Delay p95 - - - How many runs were waiting, over the selected time. marks - where the queue was throttled. - - } - > - Backlog - - - Pause/resume - - - - - {queueRows.length > 0 ? ( - queueRows.map((queue) => { - const limit = queue.concurrencyLimit ?? environment.concurrencyLimit; - const isAtConcurrencyLimit = queue.running >= limit; - const isAtQueueLimit = - environment.queueSizeLimit !== null && - queue.queued >= environment.queueSizeLimit; - const queueFilterableName = queueMetricsKey(queue); - const queueMetric = metricsByQueue[queueFilterableName]; - const queueDetailPath = v3QueuePath(organization, project, env, { - friendlyId: queue.id, - }); - return ( - - s, so - // they render beside the link (leading/trailing), never inside it — - // otherwise the cell is invalid
- - - - ); -} - -function EnvironmentPauseResumeButton({ - env, - disabled, -}: { - env: { type: RuntimeEnvironmentType; paused: boolean }; - disabled: boolean; -}) { - const navigation = useNavigation(); - const [isOpen, setIsOpen] = useState(false); - - useEffect(() => { - if (navigation.state === "loading" || navigation.state === "idle") { - // oxlint-disable-next-line react/set-state-in-effect -- This effect intentionally synchronizes route state after an external or lifecycle change. - setIsOpen(false); - } - }, [navigation.state]); - - const isLoading = Boolean( - navigation.formData?.get("action") === (env.paused ? "environment-resume" : "environment-pause") - ); - - return ( - -
- - - -
- - - -
- - - {disabled - ? "You don't have permission to manage queues" - : env.paused - ? `Resumes ${environmentFullTitle(env)} so its runs can be dequeued again.` - : `Pauses all runs from being dequeued in ${environmentFullTitle(env)}. Any executing runs will continue to run.`} - - - -
- - {env.paused ? "Resume environment?" : "Pause environment?"} -
- - {env.paused - ? `This will allow runs to be dequeued in ${environmentFullTitle(env)} again.` - : `This will pause all runs from being dequeued in ${environmentFullTitle( - env - )}. Any executing runs will continue to run.`} - -
setIsOpen(false)}> - - : env.paused ? PlayIcon : PauseIcon - } - shortcut={{ modifiers: ["mod"], key: "enter" }} - > - {env.paused ? "Resume environment" : "Pause environment"} - - } - cancelButton={ - - - - } - /> - -
- - - ); -} - -export function isEnvironmentPauseResumeFormSubmission( - formMethod: string | undefined, - formData: FormData | undefined -) { - if (!formMethod || !formData) { - return false; - } - - return ( - formMethod.toLowerCase() === "post" && - (formData.get("action") === "environment-pause" || - formData.get("action") === "environment-resume") - ); -} - -export function QueueFilters() { - return ; -} - -type MetricTileRow = Record; - -type TilePoint = { bucket: number; value: number | null }; - -// Inline colour swatch matching the chart's warning ("yellow") line — used in tooltip copy that -// refers to that colour instead of naming it, so the swatch always matches the chart. -function WarningSwatch() { - return ( - - ); -} - -type QueueHeaderTile = { - id: string; - label: string; - /** Info-icon copy explaining what the chart shows, rendered next to the card title. */ - description: ReactNode; - color: string; - /** Optional inline legend rendered below the card title: a fixed set of {colored square, label} - * entries, for charts where a colour (e.g. the orange warning line) needs explaining. */ - legend?: Array<{ color: string; label: string }>; - query: string; - /** Formats a single bucket's value in the chart tooltip. */ - formatValue?: (value: number) => string; - /** Formats the y-axis tick labels. Without it the axis shows raw numbers (bad for durations - * in ms or percent scales). Passed through to Chart.Line's yAxisProps.tickFormatter. */ - formatAxis?: (value: number) => string; - /** Hover tooltip explaining the headline readout next to the title (e.g. what "9% of current - * period" means). Without it the readout has no tooltip. */ - totalTooltip?: string; - /** Turns one row per bucket into the per-bucket points the chart draws. A null value is a - * bucket the metric has nothing to say about, and the line breaks there rather than reading 0. */ - derive: (rows: MetricTileRow[]) => { - points: TilePoint[]; - total: number; - formatTotal?: (total: number) => string; - totalClassName?: string; - }; - /** - * Optional second query, run at the range's natural bucket width, that owns the headline readout. - * For a headline that is not invariant to bucket width — a share of buckets, or a percentile, - * as opposed to a max over gauges — deriving it from the plotted rows would move the number - * whenever this tile's floor widens them. Only requested while the floor is actually widening - * anything; on ranges whose natural width is already at or above the floor the chart's own rows - * are identical and are reused. - */ - readout?: { - query: string; - derive: (rows: MetricTileRow[]) => { - total: number; - formatTotal?: (total: number) => string; - totalClassName?: string; - }; - }; -}; - -function tileNumber(value: number | string | null): number { - const n = typeof value === "number" ? value : Number(value); - return Number.isFinite(n) ? n : 0; -} - -function tileTimeToMs(value: number | string | null): number { - const s = String(value).replace(" ", "T"); - return Date.parse(s.endsWith("Z") ? s : `${s}Z`); -} - -/** Peak of a series, ignoring the buckets it has nothing to say about. */ -function peakOf(points: TilePoint[]): number { - return points.reduce((max, p) => (p.value === null ? max : Math.max(max, p.value)), 0); -} - -const SCHEDULING_DELAY_QUERY = `SELECT timeBucket() AS t,\n round(quantilesTDigestMerge(0.5, 0.9, 0.95, 0.99)(wait_quantiles)[3]) AS p95,\n sum(wait_ms_count) AS samples\nFROM env_metrics\nGROUP BY t\nORDER BY t`; - -const THROTTLED_QUERY = `SELECT timeBucket() AS t,\n sum(throttled_count) AS throttled\nFROM env_metrics\nGROUP BY t\nORDER BY t`; - -const QUEUE_HEADER_TILES: QueueHeaderTile[] = [ - { - id: "saturation", - label: "Env saturation", - description: ( - <> - How much of the environment's concurrency is in use. Turns above 100%, - when it's into burst capacity. - - ), - color: "var(--color-queues-chart)", - legend: [ - { color: "var(--color-queues-chart)", label: "Saturation" }, - { color: "var(--color-warning)", label: "Over limit" }, - ], - query: `SELECT timeBucket() AS t,\n max(max_env_running) AS running,\n max(max_env_limit) AS env_limit\nFROM env_metrics\nGROUP BY t\nORDER BY t`, - formatValue: (v) => (v > 100 ? `${v}% — over the environment limit` : `${v}%`), - formatAxis: (v) => `${v}%`, - derive: (rows) => { - const points = rows.map((r) => { - const limit = tileNumber(r.env_limit); - return { - bucket: tileTimeToMs(r.t), - value: limit > 0 ? Math.round((tileNumber(r.running) / limit) * 100) : 0, - }; - }); - return { points, total: peakOf(points), formatTotal: (v) => `${v}% peak` }; - }, - }, - { - id: "backlog", - label: "Backlog", - description: "How many runs are waiting across the environment, over time.", - color: "var(--color-queues-chart)", - query: `SELECT timeBucket() AS t,\n max(max_env_queued) AS queued\nFROM env_metrics\nGROUP BY t\nORDER BY t`, - derive: (rows) => { - const points = rows.map((r) => ({ - bucket: tileTimeToMs(r.t), - value: tileNumber(r.queued), - })); - return { points, total: peakOf(points), formatTotal: (v) => `${v.toLocaleString()} peak` }; - }, - }, - { - id: "p95", - label: "Scheduling delay p95", - description: ( - <> - How long runs wait before they start (95% start faster than this). Turns {" "} - above 1 minute. - - ), - totalTooltip: "The worst p95 in the selected window.", - color: "var(--color-queues-chart)", - legend: [ - { color: "var(--color-queues-chart)", label: "p95" }, - { color: "var(--color-warning)", label: "Over 1 min" }, - ], - query: SCHEDULING_DELAY_QUERY, - formatValue: formatWaitMs, - formatAxis: formatWaitMs, - derive: (rows) => { - const points = rows.map((r) => ({ - bucket: tileTimeToMs(r.t), - value: tileNumber(r.samples) > 0 ? tileNumber(r.p95) : null, - })); - return { points, total: peakOf(points) }; - }, - readout: { - query: SCHEDULING_DELAY_QUERY, - /** - * Merging quantile states over a wider bucket yields a p95 between the sub-buckets' own, so - * the worst p95 has to be read at the range's natural width or a burst of slow starts shorter - * than the plotted bucket is averaged away. Unlike the gauges, whose max of maxes is the same - * at any width. - */ - derive: (rows) => { - const worst = rows.reduce( - (max, r) => (tileNumber(r.samples) > 0 ? Math.max(max, tileNumber(r.p95)) : max), - 0 - ); - return { - total: worst, - formatTotal: (v) => (v > 0 ? formatWaitMs(v) : "–"), - totalClassName: worst >= 60_000 ? "text-warning" : undefined, - }; - }, - }, - }, - { - id: "throttled", - label: "Throttled", - description: "How often runs were held back by a limit.", - totalTooltip: "The share of the selected window with at least one blocked dequeue.", - color: "var(--color-queues-chart)", - legend: [{ color: "var(--color-warning)", label: "Throttled" }], - query: THROTTLED_QUERY, - derive: (rows) => { - const points = rows.map((r) => ({ - bucket: tileTimeToMs(r.t), - value: tileNumber(r.throttled), - })); - return { points, total: peakOf(points) }; - }, - readout: { - query: THROTTLED_QUERY, - /** - * Share of the window that saw any throttling. A raw event sum isn't interpretable (it - * scales with poll rate and window length); the fraction of buckets with a throttle is. - * Gap fill zero-fills this counter, so every bucket in the window is present and the row - * count is the honest denominator. - */ - derive: (rows) => { - const nonzero = rows.filter((r) => tileNumber(r.throttled) > 0).length; - const pct = rows.length > 0 ? Math.round((nonzero / rows.length) * 100) : 0; - return { - total: pct, - formatTotal: (v) => `${v}% of current period`, - totalClassName: pct > 0 ? "text-warning" : undefined, - }; - }, - }, - }, -]; - -/** - * Bucket floor shared by every hero tile. Scheduling delay and throttling are event-driven, so at - * the 10-second width a short range would otherwise pick, most buckets hold no samples at all. One - * floor for all four keeps their x-axes identical, which the shared hover crosshair relies on. - */ -const HERO_CHART_MIN_BUCKET_SECONDS = 60; - -type TileTimeRange = MetricResourceTimeRange; - -// Full-size env metric chart rendered inside a ChartCard. Same data path as before -// (client-side TRQL via useMetricResourceQuery with fillGaps), drawn as a line that -// participates in the shared hover + drag-to-zoom of the enclosing ChartSyncProvider. -function QueueEnvMetricChart({ - tile, - timeRange, - referenceLines, - thresholdStroke, - warningOverlay, - solidWarning = false, -}: { - tile: QueueHeaderTile; - timeRange: TileTimeRange; - referenceLines?: Array<{ - y: number; - label?: string; - labelPlacement?: "inside" | "outside"; - }>; - thresholdStroke?: { value: number; aboveColor: string }; - warningOverlay?: { threshold: number }; - /** When set, the ENTIRE line turns warning-coloured if the series is ever non-zero (used for - * throttling: any throttle in the window colours the whole line). Mutually exclusive with the - * per-bucket warningOverlay. */ - solidWarning?: boolean; -}) { - const organization = useOrganization(); - const project = useProject(); - const environment = useEnvironment(); - - const sharedOptions = { - organizationId: organization.id, - projectId: project.id, - environmentId: environment.id, - timeRange, - defaultPeriod: QUEUE_METRICS_DEFAULT_PERIOD, - fillGaps: true, - }; - - const { rows, showLoading, failed } = useMetricResourceQuery(tile.query, { - ...sharedOptions, - minBucketSeconds: HERO_CHART_MIN_BUCKET_SECONDS, - }); - - const derived = tile.derive(rows); - const points = derived.points; - - const plottedBucketMs = points.length > 1 ? points[1]!.bucket - points[0]!.bucket : 0; - const floorWidenedBuckets = - plottedBucketMs > 0 && plottedBucketMs <= HERO_CHART_MIN_BUCKET_SECONDS * 1000; - const readoutQuery = tile.readout && floorWidenedBuckets ? tile.readout.query : ""; - const readoutResult = useMetricResourceQuery(readoutQuery, sharedOptions); - - const { total, formatTotal, totalClassName } = tile.readout - ? tile.readout.derive(readoutQuery ? readoutResult.rows : rows) - : derived; - - // Same point shape the shared axis/tooltip helpers expect. - const data = points - .map((p) => ({ bucket: p.bucket, [tile.id]: p.value })) - .filter((p) => Number.isFinite(p.bucket)); - - // Whole-line warning colour when the series was ever non-zero (throttling: one throttle in the - // window colours the entire line). Otherwise the tile's normal colour. - const wholeLineWarning = solidWarning && total > 0; - const lineColor = wholeLineWarning ? "var(--color-warning)" : tile.color; - - const chartConfig = useMemo( - () => ({ [tile.id]: { label: tile.label, color: lineColor } }), - [tile.id, tile.label, lineColor] - ); - - const { tickFormatter, tooltipLabelFormatter } = buildActivityTimeAxis(data); - const hasData = data.length > 0 && data.some((p) => Number(p[tile.id] ?? 0) > 0); - - // Peak readout lives in the card title (ChartCard has no dedicated value slot). A zero/empty - // total renders no readout at all (skipping "0% peak", "0 peak", "0" and the p95 "–" placeholder) - // so the card title stands alone until there's a non-zero value to show. - const readoutLoading = tile.readout ? readoutResult.showLoading : showLoading; - const readoutFailed = tile.readout ? readoutResult.failed : failed; - const peak = readoutLoading ? ( - - ) : readoutFailed || total === 0 ? null : formatTotal ? ( - formatTotal(total) - ) : ( - total.toLocaleString() - ); - - return ( - - - - {tile.label} - - - {peak != null ? ( - tile.totalTooltip && !readoutLoading ? ( - - {peak} - - } - content={tile.totalTooltip} - className="max-w-[230px]" - disableHoverableContent - /> - ) : ( - - {peak} - - ) - ) : null} - - {tile.legend && (showLoading || hasData) ? ( - - {tile.legend.map((item) => ( - - - {item.label} - - ))} - - ) : null} - - } - > - {showLoading ? ( - - ) : failed ? ( -
- Unable to load metrics -
- ) : hasData ? ( - - - - ) : ( -
- No activity -
- )} - - ); -} - -function QueueMetricChartSkeleton() { - return ( -
- {Array.from({ length: 42 }).map((_, i) => ( -
- ))} -
- ); -} - -/** Health as a stock Badge: color carries the state, w-fit keeps it content-width. */ -type QueueHealth = { - paused: boolean; - running: number; - queued: number; - limit: number; -}; - -type QueueHealthLabel = "Paused" | "At capacity" | "Backlogged" | "Active" | "Idle"; - -// Single source of truth for the queue health decision, shared by the badge and the table's -// health-column sort so the sorted order always matches the labels shown. -function queueHealthLabel({ paused, running, queued, limit }: QueueHealth): QueueHealthLabel { - if (paused) return "Paused"; - if (isQueueAtCapacity({ running, queued, limit })) return "At capacity"; - if (queued > 0) return "Backlogged"; - if (running > 0) return "Active"; - return "Idle"; -} - -// Tint + colored text, sized like the error status chips (see ErrorStatusBadge). -const QUEUE_HEALTH_STYLES: Record = { - Paused: "bg-warning/10 text-warning system:bg-warning system:text-white", - "At capacity": "bg-warning/10 text-warning system:bg-warning system:text-white", - Backlogged: "bg-blue-500/10 text-blue-500 system:bg-blue-500 system:text-white", - Active: "bg-success/10 text-success system:bg-success system:text-white", - Idle: "bg-charcoal-500/10 text-text-dimmed system:bg-charcoal-500 system:text-white", -}; - -function QueueHealthBadge(health: QueueHealth) { - const label = queueHealthLabel(health); - return ( - - {label} - - ); -} - -// The `queue_metrics`-prefixed key a queue is stored under (task queues are prefixed `task/`). -function queueMetricsKey(queue: { type: string; name: string }): string { - return `${queue.type === "task" ? "task/" : ""}${queue.name}`; -} - -function formatWaitMs(ms: number): string { - if (ms < 1000) return `${Math.round(ms)}ms`; - if (ms < 60_000) return `${(ms / 1000).toFixed(1)}s`; - if (ms < 3_600_000) return `${(ms / 60_000).toFixed(1)}m`; - return `${(ms / 3_600_000).toFixed(1)}h`; -} - -// Drop a trailing ".00" from whole percentages so "50.00" reads as "50" but "12.50" is preserved. -function formatOverridePercent(percent: number): string { - return Number.isInteger(percent) ? percent.toString() : percent.toFixed(2).replace(/\.?0+$/, ""); -} - -// Classic Queues page, restored verbatim from before the Queue Metrics feature. Rendered -// when queueMetricsUiEnabled is off so a gated org sees exactly the pre-metrics UI. -function ClassicQueuesView() { - const { environment, queues, pagination, hasFilters, autoReloadPollIntervalMs, canWriteTasks } = - useTypedLoaderData(); - - const organization = useOrganization(); - const project = useProject(); - const env = useEnvironment(); - const plan = useCurrentPlan(); - - useAutoRevalidate({ interval: autoReloadPollIntervalMs, onFocus: true }); - - const { limitStatus, limitClassName } = getEnvConcurrencyLimitStatus(environment); - - return ( - - - - - - - - Queues docs - - - - - -
-
- paused : undefined} - animate - accessory={ -
- {environment.runsEnabled && - env.pauseSource !== ENVIRONMENT_PAUSE_SOURCE_BILLING_LIMIT ? ( - - ) : null} - -
- } - valueClassName={env.paused ? "text-warning tabular-nums" : "tabular-nums"} - compactThreshold={1000000} - /> - - Including {environment.running - environment.concurrencyLimit} burst runs{" "} - - - ) : limitStatus === "limit" ? ( - "At concurrency limit" - ) : undefined - } - accessory={ - - } - compactThreshold={1000000} - /> - 1 ? ( - - Burst limit {environment.burstFactor * environment.concurrencyLimit}{" "} - - - ) : undefined - } - accessory={ - plan ? ( - plan?.v3Subscription?.plan?.limits.concurrentRuns.canExceed ? ( - - Increase limit - - ) : ( - - Increase limit - - ) - ) : null - } - /> -
- -
-
- - -
- - - - Name - Queued - Running - Limit - -
- Environment - - This queue is limited by your environment's concurrency limit of{" "} - {environment.concurrencyLimit}. - -
-
- User - - This queue is limited by a concurrency limit set in your code. - -
-
- Override - - This queue's concurrency limit has been manually overridden from the - dashboard or API. - -
- - } - > - Limited by - - - Pause/resume - - - - - {queues.length > 0 ? ( - queues.map((queue) => { - const limit = queue.concurrencyLimit ?? environment.concurrencyLimit; - const isAtConcurrencyLimit = queue.running >= limit; - const isAtQueueLimit = - environment.queueSizeLimit !== null && - queue.queued >= environment.queueSizeLimit; - const queueFilterableName = `${queue.type === "task" ? "task/" : ""}${ - queue.name - }`; - return ( - - - - - {queue.concurrency?.overriddenAt ? ( - - Concurrency limit overridden - - } - content="This queue's concurrency limit has been manually overridden from the dashboard or API." - className="max-w-[230px]" - disableHoverableContent - /> - ) : null} - {queue.paused ? ( - - Paused - - ) : null} - {isAtQueueLimit ? ( - - At queue limit - - ) : null} - {isAtConcurrencyLimit ? ( - - At concurrency limit - - ) : null} - - - - {queue.queued} - - 0 && "text-text-bright", - isAtConcurrencyLimit && "text-warning" - )} - > - {queue.running} - - - {limit} - - - {queue.concurrency?.overriddenAt ? ( - Override - ) : queue.concurrencyLimit ? ( - "User" - ) : ( - "Environment" - )} - - - ) - } - hiddenButtons={ - !queue.paused && ( - - ) - } - popoverContent={ - <> - {queue.paused ? ( - - ) : ( - - )} - - - - - - - } - /> - - ); - }) - ) : ( - - -
- - {hasFilters ? "No queues found matching your filters" : "No queues found"} - -
- - - )} - -
-
-
- - - ); -} - -function BurstFactorTooltip({ - environment, -}: { - environment: { burstFactor: number; concurrencyLimit: number }; -}) { - return ( - - ); -} diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.tasks.scheduled.$taskParam/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.tasks.scheduled.$taskParam/route.tsx index 446d6a8b260..d51d6fdd28e 100644 --- a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.tasks.scheduled.$taskParam/route.tsx +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.tasks.scheduled.$taskParam/route.tsx @@ -98,7 +98,7 @@ import { v3EditSchedulePath, v3EnvironmentPath, v3NewSchedulePath, - v3QueuePath, + concurrencyQueuePath, v3RunsPath, v3SchedulePath, v3SchedulesAddOnPath, @@ -266,7 +266,9 @@ export default function Page() { taskIdentifier: task.slug, }); const queuePath = task.queue - ? v3QueuePath(organization, project, environment, { friendlyId: task.queue.friendlyId }) + ? concurrencyQueuePath(organization, project, environment, { + friendlyId: task.queue.friendlyId, + }) : undefined; const filters: TaskRunListSearchFilters = useMemo(() => ({ tasks: [task.slug] }), [task.slug]); diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.tasks.standard.$taskParam/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.tasks.standard.$taskParam/route.tsx index 56f553547b1..6615a2fc8d6 100644 --- a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.tasks.standard.$taskParam/route.tsx +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.tasks.standard.$taskParam/route.tsx @@ -63,8 +63,8 @@ import { requireUser } from "~/services/session.server"; import { EnvironmentParamSchema, v3EnvironmentPath, - v3QueuePath, - v3QueuesPath, + concurrencyQueuePath, + concurrencyPath, v3TestTaskPath, } from "~/utils/pathBuilder"; import { parseFiniteInt } from "~/utils/searchParams"; @@ -198,9 +198,11 @@ export default function Page() { const testPath = v3TestTaskPath(organization, project, environment, { taskIdentifier: task.slug, }); - const queuesPath = v3QueuesPath(organization, project, environment); + const queuesPath = concurrencyPath(organization, project, environment); const queuePath = task.queue - ? v3QueuePath(organization, project, environment, { friendlyId: task.queue.friendlyId }) + ? concurrencyQueuePath(organization, project, environment, { + friendlyId: task.queue.friendlyId, + }) : undefined; const { value } = useSearchParams(); @@ -527,7 +529,7 @@ function TaskActivityCard({ > {view === "queue" ? ( ({ type: "queues" }), + }, + corsStrategy: "all", + }, + async ({ params, body, authentication }) => { + return concurrencyLimitsSystem.limits + .override(authentication.environment, params.name, body) + .match( + (limit) => json(limit), + (error) => { + switch (error.type) { + case "limit_not_found": + return json({ error: "Concurrency limit not found" }, { status: 404 }); + case "conflict": + return json( + { error: "The limit changed concurrently; retry the request" }, + { status: 409 } + ); + case "invalid_override": + return json({ error: error.message }, { status: 400 }); + default: + return json({ error: "Failed to override concurrency limit" }, { status: 500 }); + } + } + ); + } +); + +export const action = route.action; +export const loader = route.loader; diff --git a/apps/webapp/app/routes/api.v1.concurrency-limits.$name.reset.ts b/apps/webapp/app/routes/api.v1.concurrency-limits.$name.reset.ts new file mode 100644 index 00000000000..5e0af9192b1 --- /dev/null +++ b/apps/webapp/app/routes/api.v1.concurrency-limits.$name.reset.ts @@ -0,0 +1,42 @@ +import { json } from "@remix-run/server-runtime"; +import { z } from "zod"; +import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server"; +import { concurrencyLimitsSystem } from "~/v3/services/concurrencyLimitsSystemInstance.server"; + +const ParamsSchema = z.object({ + name: z.string(), +}); + +const route = createActionApiRoute( + { + params: ParamsSchema, + authorization: { + action: "write", + resource: () => ({ type: "queues" }), + }, + corsStrategy: "all", + }, + async ({ params, authentication }) => { + return concurrencyLimitsSystem.limits.reset(authentication.environment, params.name).match( + (limit) => json(limit), + (error) => { + switch (error.type) { + case "limit_not_found": + return json({ error: "Concurrency limit not found" }, { status: 404 }); + case "conflict": + return json( + { error: "The limit changed concurrently; retry the request" }, + { status: 409 } + ); + case "limit_not_overridden": + return json({ error: "Concurrency limit has no override to reset" }, { status: 400 }); + default: + return json({ error: "Failed to reset concurrency limit" }, { status: 500 }); + } + } + ); + } +); + +export const action = route.action; +export const loader = route.loader; diff --git a/apps/webapp/app/routes/api.v1.concurrency-limits.$name.ts b/apps/webapp/app/routes/api.v1.concurrency-limits.$name.ts new file mode 100644 index 00000000000..3035de12baf --- /dev/null +++ b/apps/webapp/app/routes/api.v1.concurrency-limits.$name.ts @@ -0,0 +1,31 @@ +import { json } from "@remix-run/server-runtime"; +import { z } from "zod"; +import { createLoaderApiRoute } from "~/services/routeBuilders/apiBuilder.server"; +import { concurrencyLimitsSystem } from "~/v3/services/concurrencyLimitsSystemInstance.server"; + +const ParamsSchema = z.object({ + name: z.string(), +}); + +export const loader = createLoaderApiRoute( + { + params: ParamsSchema, + findResource: async () => 1, + authorization: { + action: "read", + resource: () => ({ type: "queues" }), + }, + corsStrategy: "all", + }, + async ({ params, authentication }) => { + return concurrencyLimitsSystem.limits.retrieve(authentication.environment, params.name).match( + (limit) => json(limit), + (error) => { + if (error.type === "limit_not_found") { + return json({ error: "Concurrency limit not found" }, { status: 404 }); + } + return json({ error: "Failed to retrieve concurrency limit" }, { status: 500 }); + } + ); + } +); diff --git a/apps/webapp/app/routes/api.v1.concurrency-limits.ts b/apps/webapp/app/routes/api.v1.concurrency-limits.ts new file mode 100644 index 00000000000..0b9c936019d --- /dev/null +++ b/apps/webapp/app/routes/api.v1.concurrency-limits.ts @@ -0,0 +1,43 @@ +import { json } from "@remix-run/server-runtime"; +import { z } from "zod"; +import { createLoaderApiRoute } from "~/services/routeBuilders/apiBuilder.server"; +import { concurrencyLimitsSystem } from "~/v3/services/concurrencyLimitsSystemInstance.server"; + +const SearchParamsSchema = z.object({ + page: z.coerce.number().int().positive().default(1), + perPage: z.coerce.number().int().positive().max(100).default(25), +}); + +export const loader = createLoaderApiRoute( + { + searchParams: SearchParamsSchema, + findResource: async () => 1, + authorization: { + action: "read", + resource: () => ({ type: "queues" }), + }, + corsStrategy: "all", + }, + async ({ searchParams, authentication }) => { + const [data, count] = await Promise.all([ + concurrencyLimitsSystem.limits.list(authentication.environment, { + page: searchParams.page, + perPage: searchParams.perPage, + }), + concurrencyLimitsSystem.limits.totalCount(authentication.environment), + ]); + + if (data.isErr() || count.isErr()) { + return json({ error: "Failed to list concurrency limits" }, { status: 500 }); + } + + return json({ + data: data.value, + pagination: { + currentPage: searchParams.page, + totalPages: Math.ceil(count.value / searchParams.perPage), + count: count.value, + }, + }); + } +); diff --git a/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.override.ts b/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.override.ts index 90f5772c5d3..4aa343ff377 100644 --- a/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.override.ts +++ b/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.override.ts @@ -1,7 +1,7 @@ import { json } from "@remix-run/server-runtime"; import { type RetrieveQueueParam, RetrieveQueueType } from "@trigger.dev/core/v3"; import { z } from "zod"; -import { toQueueItem } from "~/presenters/v3/QueueRetrievePresenter.server"; +import { toPublicQueueItem, toQueueItem } from "~/presenters/v3/QueueRetrievePresenter.server"; import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server"; import { concurrencySystem } from "~/v3/services/concurrencySystemInstance.server"; import { @@ -46,27 +46,41 @@ const route = createActionApiRoute( body.percent !== undefined ? { percent: body.percent } : { limit: body.concurrencyLimit! }; return concurrencySystem.queues - .overrideQueueConcurrencyLimit(authentication.environment, input, override) + .overrideQueueConcurrencyLimit(authentication.environment, input, override, undefined, { + v1Only: true, + }) .match( (queue) => { return json( - toQueueItem({ - friendlyId: queue.friendlyId, - name: queue.name, - type: queue.type, - running: queue.running, - queued: queue.queued, - concurrencyLimit: queue.concurrencyLimit, - concurrencyLimitBase: queue.concurrencyLimitBase, - concurrencyLimitOverriddenAt: queue.concurrencyLimitOverriddenAt, - concurrencyLimitOverriddenBy: null, - paused: queue.paused, - }), + toPublicQueueItem( + toQueueItem({ + friendlyId: queue.friendlyId, + name: queue.name, + type: queue.type, + version: queue.concurrencyVersion, + running: queue.running, + queued: queue.queued, + concurrencyLimit: queue.concurrencyLimit, + concurrencyLimitBase: queue.concurrencyLimitBase, + concurrencyLimitOverriddenAt: queue.concurrencyLimitOverriddenAt, + concurrencyLimitOverriddenBy: null, + paused: queue.paused, + }) + ), { status: 200 } ); }, (error) => { switch (error.type) { + case "queue_version_unsupported": { + return json( + { + error: + "This queue's concurrency is declared with the task `concurrency` option; manage it through the concurrency-limits endpoints instead (a task's inline limit lives under its derived `task/` name)", + }, + { status: 400 } + ); + } case "queue_not_found": { return json({ error: "Queue not found" }, { status: 404 }); } @@ -74,6 +88,9 @@ const route = createActionApiRoute( case "concurrency_limit_exceeds_maximum": { return json({ error: error.message }, { status: 400 }); } + case "concurrent_modification": { + return json({ error: error.message }, { status: 409 }); + } case "queue_update_failed": { return json({ error: "Failed to update queue concurrency limit" }, { status: 500 }); } diff --git a/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.reset.ts b/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.reset.ts index 503d875e471..3abb2a5c52c 100644 --- a/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.reset.ts +++ b/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.reset.ts @@ -1,7 +1,7 @@ import { json } from "@remix-run/server-runtime"; import { type RetrieveQueueParam, RetrieveQueueType } from "@trigger.dev/core/v3"; import { z } from "zod"; -import { toQueueItem } from "~/presenters/v3/QueueRetrievePresenter.server"; +import { toPublicQueueItem, toQueueItem } from "~/presenters/v3/QueueRetrievePresenter.server"; import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server"; import { concurrencySystem } from "~/v3/services/concurrencySystemInstance.server"; @@ -29,52 +29,69 @@ const route = createActionApiRoute( name: decodeURIComponent(params.queueParam).replace(/%2F/g, "/"), }; - return concurrencySystem.queues.resetConcurrencyLimit(authentication.environment, input).match( - (queue) => { - return json( - toQueueItem({ - friendlyId: queue.friendlyId, - name: queue.name, - type: queue.type, - running: queue.running, - queued: queue.queued, - concurrencyLimit: queue.concurrencyLimit, - concurrencyLimitBase: queue.concurrencyLimitBase, - concurrencyLimitOverriddenAt: queue.concurrencyLimitOverriddenAt, - concurrencyLimitOverriddenBy: null, - paused: queue.paused, - }), - { status: 200 } - ); - }, - (error) => { - switch (error.type) { - case "queue_not_found": { - return json({ error: "Queue not found" }, { status: 404 }); - } - case "queue_not_overridden": { - return json({ error: "Queue is not overridden" }, { status: 400 }); - } - case "queue_update_failed": { - return json({ error: "Failed to update queue concurrency limit" }, { status: 500 }); - } - case "sync_queue_concurrency_to_engine_failed": { - return json( - { error: "Failed to sync queue concurrency limit to engine" }, - { status: 500 } - ); - } - case "get_queue_stats_failed": { - return json({ error: "Failed to get queue stats" }, { status: 500 }); - } - case "other": - default: { - error.type satisfies "other"; - return json({ error: "Internal server error" }, { status: 500 }); + return concurrencySystem.queues + .resetConcurrencyLimit(authentication.environment, input, { v1Only: true }) + .match( + (queue) => { + return json( + toPublicQueueItem( + toQueueItem({ + friendlyId: queue.friendlyId, + name: queue.name, + type: queue.type, + version: queue.concurrencyVersion, + running: queue.running, + queued: queue.queued, + concurrencyLimit: queue.concurrencyLimit, + concurrencyLimitBase: queue.concurrencyLimitBase, + concurrencyLimitOverriddenAt: queue.concurrencyLimitOverriddenAt, + concurrencyLimitOverriddenBy: null, + paused: queue.paused, + }) + ), + { status: 200 } + ); + }, + (error) => { + switch (error.type) { + case "queue_version_unsupported": { + return json( + { + error: + "This queue's concurrency is declared with the task `concurrency` option; manage it through the concurrency-limits endpoints instead (a task's inline limit lives under its derived `task/` name)", + }, + { status: 400 } + ); + } + case "queue_not_found": { + return json({ error: "Queue not found" }, { status: 404 }); + } + case "queue_not_overridden": { + return json({ error: "Queue is not overridden" }, { status: 400 }); + } + case "concurrent_modification": { + return json({ error: error.message }, { status: 409 }); + } + case "queue_update_failed": { + return json({ error: "Failed to update queue concurrency limit" }, { status: 500 }); + } + case "sync_queue_concurrency_to_engine_failed": { + return json( + { error: "Failed to sync queue concurrency limit to engine" }, + { status: 500 } + ); + } + case "get_queue_stats_failed": { + return json({ error: "Failed to get queue stats" }, { status: 500 }); + } + case "other": + default: { + error.type satisfies "other"; + return json({ error: "Internal server error" }, { status: 500 }); + } } } - } - ); + ); } ); diff --git a/apps/webapp/app/routes/api.v1.queues.$queueParam.metrics.ts b/apps/webapp/app/routes/api.v1.queues.$queueParam.metrics.ts index f4d1e8cf8ae..838b2b904d2 100644 --- a/apps/webapp/app/routes/api.v1.queues.$queueParam.metrics.ts +++ b/apps/webapp/app/routes/api.v1.queues.$queueParam.metrics.ts @@ -51,7 +51,7 @@ export const loader = createLoaderApiRoute( findResource: async () => 1, // dummy — the queue name isn't resolved against Postgres authorization: { action: "read", - resource: () => ({ type: "query", id: "queue_metrics" }), + resource: () => ({ type: "query", id: "concurrency_metrics" }), }, }, async ({ params, searchParams, authentication }) => { diff --git a/apps/webapp/app/routes/api.v1.queues.$queueParam.pause.ts b/apps/webapp/app/routes/api.v1.queues.$queueParam.pause.ts index dd5f43da5d1..a51777fefb2 100644 --- a/apps/webapp/app/routes/api.v1.queues.$queueParam.pause.ts +++ b/apps/webapp/app/routes/api.v1.queues.$queueParam.pause.ts @@ -1,5 +1,6 @@ import { json } from "@remix-run/server-runtime"; -import { type QueueItem, type RetrieveQueueParam, RetrieveQueueType } from "@trigger.dev/core/v3"; +import { type RetrieveQueueParam, RetrieveQueueType } from "@trigger.dev/core/v3"; +import { toPublicQueueItem } from "~/presenters/v3/QueueRetrievePresenter.server"; import { z } from "zod"; import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server"; import { PauseQueueService } from "~/v3/services/pauseQueue.server"; @@ -44,8 +45,7 @@ const route = createActionApiRoute( return json({ error: result.code }, { status: 400 }); } - const q: QueueItem = result.queue; - return json(q); + return json(toPublicQueueItem(result.queue)); } ); diff --git a/apps/webapp/app/routes/api.v1.queues.$queueParam.ts b/apps/webapp/app/routes/api.v1.queues.$queueParam.ts index 434046d961d..82503142903 100644 --- a/apps/webapp/app/routes/api.v1.queues.$queueParam.ts +++ b/apps/webapp/app/routes/api.v1.queues.$queueParam.ts @@ -1,7 +1,10 @@ import { json } from "@remix-run/server-runtime"; -import { type QueueItem, type RetrieveQueueParam, RetrieveQueueType } from "@trigger.dev/core/v3"; +import { type RetrieveQueueParam, RetrieveQueueType } from "@trigger.dev/core/v3"; import { z } from "zod"; -import { QueueRetrievePresenter } from "~/presenters/v3/QueueRetrievePresenter.server"; +import { + QueueRetrievePresenter, + toPublicQueueItem, +} from "~/presenters/v3/QueueRetrievePresenter.server"; import { createLoaderApiRoute } from "~/services/routeBuilders/apiBuilder.server"; import { determineEngineVersion } from "~/v3/engineVersion.server"; @@ -52,7 +55,6 @@ export const loader = createLoaderApiRoute( return json({ error: result.code }, { status: 404 }); } - const q: QueueItem = result.queue; - return json(q); + return json(toPublicQueueItem(result.queue)); } ); diff --git a/apps/webapp/app/routes/api.v1.queues.ts b/apps/webapp/app/routes/api.v1.queues.ts index 43582b9737f..fbe5da48a55 100644 --- a/apps/webapp/app/routes/api.v1.queues.ts +++ b/apps/webapp/app/routes/api.v1.queues.ts @@ -1,10 +1,10 @@ import { json } from "@remix-run/server-runtime"; -import { type QueueItem } from "@trigger.dev/core/v3"; import { z } from "zod"; import { QUEUE_LIST_DEFAULT_ITEMS_PER_PAGE, QueueListPresenter, } from "~/presenters/v3/QueueListPresenter.server"; +import { toPublicQueueItem } from "~/presenters/v3/QueueRetrievePresenter.server"; import { toOffsetLimitQueueListPagination } from "~/presenters/v3/queueListPagination.server"; import { logger } from "~/services/logger.server"; import { createLoaderApiRoute } from "~/services/routeBuilders/apiBuilder.server"; @@ -45,7 +45,7 @@ export const loader = createLoaderApiRoute( page: searchParams.page ?? 1, }); - const queues: QueueItem[] = result.queues; + const queues = result.queues.map(toPublicQueueItem); return json( { data: queues, diff --git a/apps/webapp/app/routes/orgs.$organizationSlug.projects.$projectParam.concurrency.ts b/apps/webapp/app/routes/orgs.$organizationSlug.projects.$projectParam.concurrency.ts index 4bde0ccaef3..9459382e62b 100644 --- a/apps/webapp/app/routes/orgs.$organizationSlug.projects.$projectParam.concurrency.ts +++ b/apps/webapp/app/routes/orgs.$organizationSlug.projects.$projectParam.concurrency.ts @@ -2,7 +2,7 @@ import { redirect, type LoaderFunctionArgs } from "@remix-run/server-runtime"; import { prisma } from "~/db.server"; import { SelectBestEnvironmentPresenter } from "~/presenters/SelectBestEnvironmentPresenter.server"; import { requireUser } from "~/services/session.server"; -import { ProjectParamSchema, v3QueuesPath } from "~/utils/pathBuilder"; +import { ProjectParamSchema, concurrencyPath } from "~/utils/pathBuilder"; export const loader = async ({ request, params }: LoaderFunctionArgs) => { const user = await requireUser(request); @@ -41,5 +41,5 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { const selector = new SelectBestEnvironmentPresenter(); const environment = await selector.selectBestEnvironment(project.id, user, project.environments); - return redirect(v3QueuesPath({ slug: organizationSlug }, project, environment)); + return redirect(concurrencyPath({ slug: organizationSlug }, project, environment)); }; diff --git a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.spans.$spanParam/route.tsx b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.spans.$spanParam/route.tsx index 0ec6698cad4..cc27d3c8e26 100644 --- a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.spans.$spanParam/route.tsx +++ b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.runs.$runParam.spans.$spanParam/route.tsx @@ -115,7 +115,7 @@ import { docsPath, v3BatchPath, v3DeploymentVersionPath, - v3QueuePath, + concurrencyQueuePath, v3RunDownloadLogsPath, v3RunIdempotencyKeyResetPath, v3RunPath, @@ -437,7 +437,7 @@ function RunBody({ const resetFetcher = useTypedFetcher(); const queuePath = queueMetrics?.queueFriendlyId - ? v3QueuePath(organization, project, environment, { + ? concurrencyQueuePath(organization, project, environment, { friendlyId: queueMetrics.queueFriendlyId, }) : undefined; @@ -1366,7 +1366,7 @@ function WaitingInQueueBlock({ responseReceivedAt, lastSuccessfulResponseAt, } = useQueueMetric( - `SELECT timeBucket() AS t, max(max_running) AS running, max(max_queued) AS queued, max(max_limit) AS q_limit\nFROM queue_metrics\nGROUP BY t\nORDER BY t`, + `SELECT timeBucket() AS t, max(max_running) AS running, max(max_queued) AS queued, max(max_limit) AS q_limit\nFROM concurrency_metrics\nGROUP BY t\nORDER BY t`, { ids: waiting.ids, timeRange: { period: "15m", from: null, to: null }, @@ -1465,7 +1465,7 @@ function WaitingInQueueBlock({
+ gate.queue.startsWith("limit/task/") + ); + const concurrencyGates = triggerLimits + ? [ + ...inlineTaskGates, + ...triggerLimits.map((name): { queue: string; concurrencyKey?: string } => ({ + queue: `limit/${name}`, + })), + ] + : undefined; + + /** + * The raw gates option replaces stored gates the same way concurrency does, so + * it also carries the inline gate over; a replay resending the stored gates + * collapses back to the original set through the dedupe below. + */ + const rawGates = request.body.options?.gates; + const requestedGates = + concurrencyGates ?? + (rawGates ? [...inlineTaskGates, ...rawGates] : undefined) ?? + taskGates ?? + undefined; + + const seenGates = new Set(); + const gates = requestedGates?.flatMap((gate) => { + const sanitized = sanitizeQueueName(gate.queue); + if (!sanitized) { + return []; + } + const dedupeKey = `${sanitized}${gate.concurrencyKey ?? ""}`; + if (seenGates.has(dedupeKey)) { + return []; + } + seenGates.add(dedupeKey); + return [{ queue: sanitized, concurrencyKey: gate.concurrencyKey }]; + }); + + /** + * Unreachable through the public schemas (three requested gates plus one inline + * gate is the ceiling), kept as a backstop so an overflowing set can never be + * silently truncated downstream. Replays of three-gate runs against a task that + * later gained an inline limit resolve to four and stay valid. + */ + if (gates && gates.length > 4) { + throw new ServiceValidationError( + `A run can hold at most four gates; this request resolves to ${gates.length}.` + ); + } + 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 +322,7 @@ export class DefaultQueueManager implements QueueManager { queueName: overriddenQueueName, taskTtl: meta?.ttl ?? undefined, taskKind: meta?.triggerSource, + taskGates: meta?.gates, }; } @@ -259,10 +339,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 +410,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 +431,7 @@ export class DefaultQueueManager implements QueueManager { select: { ttl: true, triggerSource: true, + gates: true, queue: { select: { id: true, name: true } }, }, }); @@ -378,6 +470,7 @@ export class DefaultQueueManager implements QueueManager { select: { ttl: true, triggerSource: true, + gates: true, queue: { select: { id: true, name: true } }, }, }); @@ -395,6 +488,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/realtime/sessionRunManager.server.ts b/apps/webapp/app/services/realtime/sessionRunManager.server.ts index a8e59335dd5..f0aeaa5bab4 100644 --- a/apps/webapp/app/services/realtime/sessionRunManager.server.ts +++ b/apps/webapp/app/services/realtime/sessionRunManager.server.ts @@ -304,6 +304,8 @@ export function buildSessionRunOptions(config: SessionTriggerConfig) { return { ...(config.machine ? { machine: config.machine as never } : {}), ...(config.queue ? { queue: { name: config.queue } } : {}), + ...(config.concurrency ? { concurrency: config.concurrency } : {}), + ...(config.concurrencyKey !== undefined ? { concurrencyKey: config.concurrencyKey } : {}), ...(config.tags ? { tags: config.tags } : {}), ...(config.maxAttempts !== undefined ? { maxAttempts: config.maxAttempts } : {}), ...(config.maxDuration !== undefined ? { maxDuration: config.maxDuration } : {}), diff --git a/apps/webapp/app/services/resolveTriggerUri.server.ts b/apps/webapp/app/services/resolveTriggerUri.server.ts index 1fe94bb6d35..98da896b5bf 100644 --- a/apps/webapp/app/services/resolveTriggerUri.server.ts +++ b/apps/webapp/app/services/resolveTriggerUri.server.ts @@ -10,7 +10,7 @@ import { import { v3DeploymentVersionPath, v3ErrorPath, - v3QueuesPath, + concurrencyPath, v3RunPath, v3RunSpanPath, v3RunsPath, @@ -127,7 +127,7 @@ function resolveInScope( // resolves to the queues list filtered to the name. return { label: parsed.name, - url: `${v3QueuesPath(organization, project, environment)}?query=${encodeURIComponent( + url: `${concurrencyPath(organization, project, environment)}?query=${encodeURIComponent( parsed.name )}`, }; 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/utils/deeplinkPages.test.ts b/apps/webapp/app/utils/deeplinkPages.test.ts index 597135cb311..04998b572fd 100644 --- a/apps/webapp/app/utils/deeplinkPages.test.ts +++ b/apps/webapp/app/utils/deeplinkPages.test.ts @@ -20,7 +20,7 @@ const ROUTES_DIR = join(APP_DIR, "routes"); const ENV_ROUTE_PREFIX = "_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam."; // Route files that name no deeplink: the environment root, and Remix's layout-opt-out spelling. -const NOT_DEEPLINK_NAMES = new Set(["_index", "queues_"]); +const NOT_DEEPLINK_NAMES = new Set(["_index", "concurrency_", "queues"]); const PROBE = "probe_01ABC"; @@ -54,7 +54,7 @@ function isRouteModule(entry: string): boolean { return existsSync(join(path, "route.tsx")) || existsSync(join(path, "route.ts")); } -// A trailing `_` only opts out of the parent layout: `queues_.$queueParam` serves `/queues/{id}`. +// A trailing `_` only opts out of the parent layout: `concurrency_.$queueParam` serves `/concurrency/{id}`. const envRoutes: string[][] = routeEntries .filter((entry) => entry.startsWith(ENV_ROUTE_PREFIX) && isRouteModule(entry)) .map((entry) => @@ -281,9 +281,9 @@ describe("the route Remix compiles from the filename", () => { expect(compiledUrl("routes/login.magic")).toBe("/login/magic"); expect( compiledUrl( - "routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.queues_.$queueParam" + "routes/_app.orgs.$organizationSlug.projects.$projectParam.env.$envParam.concurrency_.$queueParam" ) - ).toBe("/orgs/:organizationSlug/projects/:projectParam/env/:envParam/queues/:queueParam"); + ).toBe("/orgs/:organizationSlug/projects/:projectParam/env/:envParam/concurrency/:queueParam"); }); }); diff --git a/apps/webapp/app/utils/deeplinkPages.ts b/apps/webapp/app/utils/deeplinkPages.ts index 1ba60b354f9..bea7bbc919c 100644 --- a/apps/webapp/app/utils/deeplinkPages.ts +++ b/apps/webapp/app/utils/deeplinkPages.ts @@ -15,6 +15,7 @@ export const ENV_PAGE_TARGETS: ReadonlyMap = new Map([ ["branches", page("branches")], ["bulk-actions", page("bulk-actions")], ["concurrency", page("concurrency")], + ["concurrency-limits", page("concurrency-limits")], ["dashboards", page("dashboards")], ["deployments", page("deployments")], ["dev-branches", page("dev-branches")], @@ -26,7 +27,7 @@ export const ENV_PAGE_TARGETS: ReadonlyMap = new Map([ ["playground", page("playground")], ["prompts", page("prompts")], ["query", page("query")], - ["queues", page("queues")], + ["queues", page("concurrency")], ["regions", page("regions")], ["runs", page("runs")], ["schedules", page("schedules")], diff --git a/apps/webapp/app/utils/pageSwitching.test.ts b/apps/webapp/app/utils/pageSwitching.test.ts index 0f8ddce84af..2782a649fc6 100644 --- a/apps/webapp/app/utils/pageSwitching.test.ts +++ b/apps/webapp/app/utils/pageSwitching.test.ts @@ -393,7 +393,7 @@ describe("pages named after a resource", () => { it("truncate to the list they were reached from", () => { expect(projectPortablePage("runs/run_123")).toBe("runs"); expect(projectPortablePage("batches/batch_123")).toBe("batches"); - expect(projectPortablePage("queues/my-queue")).toBe("queues"); + expect(projectPortablePage("concurrency/my-queue")).toBe("concurrency"); expect(projectPortablePage("schedules/sched_123")).toBe("schedules"); expect(projectPortablePage("schedules/edit/sched_123")).toBe("schedules"); expect(projectPortablePage("deployments/deploy_123")).toBe("deployments"); @@ -677,11 +677,11 @@ describe("pathForEnvironmentSwitch", () => { expect( pathForEnvironmentSwitch({ - location: locationOn("queues/my-queue", "?page=2"), + location: locationOn("concurrency/my-queue", "?page=2"), environmentPathname: environmentLocation.pathname, environmentSlug: "prod", }) - ).toBe("/orgs/acme/projects/api/env/prod/queues"); + ).toBe("/orgs/acme/projects/api/env/prod/concurrency"); }); it("only swaps the environment slug when it cannot tell where the environment path ends", () => { diff --git a/apps/webapp/app/utils/pathBuilder.ts b/apps/webapp/app/utils/pathBuilder.ts index 38e4b9c7505..1f57fb1815f 100644 --- a/apps/webapp/app/utils/pathBuilder.ts +++ b/apps/webapp/app/utils/pathBuilder.ts @@ -588,21 +588,21 @@ export function v3SchedulesAddOnPath(organization: OrgForPath) { return `/resources/orgs/${organizationParam(organization)}/schedules-addon`; } -export function v3QueuesPath( +export function concurrencyPath( organization: OrgForPath, project: ProjectForPath, environment: EnvironmentForPath ) { - return `${v3EnvironmentPath(organization, project, environment)}/queues`; + return `${v3EnvironmentPath(organization, project, environment)}/concurrency`; } -export function v3QueuePath( +export function concurrencyQueuePath( organization: OrgForPath, project: ProjectForPath, environment: EnvironmentForPath, queue: { friendlyId: string } ) { - return `${v3QueuesPath(organization, project, environment)}/${queue.friendlyId}`; + return `${concurrencyPath(organization, project, environment)}/${queue.friendlyId}`; } export function v3WaitpointTokensPath( @@ -807,12 +807,12 @@ export function branchesDevPath( return `${v3EnvironmentPath(organization, project, environment)}/dev-branches`; } -export function concurrencyPath( +export function concurrencyLimitsPath( organization: OrgForPath, project: ProjectForPath, environment: EnvironmentForPath ) { - return `${v3EnvironmentPath(organization, project, environment)}/concurrency`; + return `${v3EnvironmentPath(organization, project, environment)}/concurrency-limits`; } export function limitsPath( diff --git a/apps/webapp/app/v3/querySchemas.ts b/apps/webapp/app/v3/querySchemas.ts index 690bbaf5396..419785a6e1e 100644 --- a/apps/webapp/app/v3/querySchemas.ts +++ b/apps/webapp/app/v3/querySchemas.ts @@ -610,12 +610,12 @@ const metricsSchema: TableSchema = { }; /** - * Schema definition for the queue_metrics table (trigger_dev.queue_metrics_v1). + * Schema definition for the concurrency_metrics table (trigger_dev.queue_metrics_v1). * Pre-aggregated into 10-second buckets. Counter columns re-aggregate with sum(), * gauges with max(), and wait_quantiles with quantilesMerge() — never FINAL. */ const queueMetricsSchema: TableSchema = { - name: "queue_metrics", + name: "concurrency_metrics", clickhouseName: "trigger_dev.queue_metrics_v1", description: "Per-queue depth, concurrency, throttling, and scheduling-delay metrics", timeConstraint: "bucket_start", @@ -770,6 +770,22 @@ const queueMetricsSchema: TableSchema = { fillMode: "carry", }), }, + max_total_running: { + name: "max_total_running", + ...column("UInt32", { + description: + "Peak in-flight runs across ALL concurrency keys of the queue in the bucket (only emitted for keyed queues). Aggregate with max().", + fillMode: "carry", + }), + }, + max_total_limit: { + name: "max_total_limit", + ...column("UInt32", { + description: + "The queue's combined concurrency limit across all keys, as stored (0 = no cap; clamp against max_env_limit). Aggregate with max().", + fillMode: "carry", + }), + }, max_ck_backlogged: { name: "max_ck_backlogged", ...column("UInt32", { @@ -828,7 +844,7 @@ const queueMetricsSchema: TableSchema = { /** * Schema definition for the env_metrics table (trigger_dev.env_metrics_v1). - * Environment-level rollup of queue_metrics with the queue dimension dropped, so + * Environment-level rollup of concurrency_metrics with the queue dimension dropped, so * header tiles and saturation charts cost the same regardless of how many queues * the environment has. Keeps the full 10-second granularity: row count is * queue-independent, so even 30-day ranges stay small. @@ -1304,7 +1320,7 @@ const llmModelsSchema: TableSchema = { * only when that key had events, so key cardinality cannot inflate the table. */ const queueMetricsByKeySchema: TableSchema = { - name: "queue_metrics_by_key", + name: "concurrency_metrics_by_key", clickhouseName: "trigger_dev.queue_metrics_ck_v1", description: "Per-concurrency-key queue metrics: backlog, throughput, and wait by key", hidden: true, @@ -1406,6 +1422,14 @@ const queueMetricsByKeySchema: TableSchema = { fillMode: "carry", }), }, + max_limit: { + name: "max_limit", + ...column("UInt32", { + description: + "The queue concurrency limit that applied to this key in the bucket (1000000 = no explicit limit). Aggregate with max().", + fillMode: "carry", + }), + }, wait_ms_sum: { name: "wait_ms_sum", ...column("UInt64", { diff --git a/apps/webapp/app/v3/queueMetricsMapping.ts b/apps/webapp/app/v3/queueMetricsMapping.ts index 9433b361a88..d341f4a63cd 100644 --- a/apps/webapp/app/v3/queueMetricsMapping.ts +++ b/apps/webapp/app/v3/queueMetricsMapping.ts @@ -131,6 +131,8 @@ export function mapEntryToRows( throttled: num(f.thr), ck_backlogged: num(f.ckq), ck_max_wait_ms: num(f.ckw), + total_running: num(f.tcc), + total_limit: num(f.tlim), }, ]; } diff --git a/apps/webapp/app/v3/runEngine.server.ts b/apps/webapp/app/v3/runEngine.server.ts index 76ac35f349b..caeb3662933 100644 --- a/apps/webapp/app/v3/runEngine.server.ts +++ b/apps/webapp/app/v3/runEngine.server.ts @@ -62,6 +62,8 @@ function createRunEngine() { queue: { defaultEnvConcurrency: env.DEFAULT_ENV_EXECUTION_CONCURRENCY_LIMIT, defaultEnvConcurrencyBurstFactor: env.DEFAULT_ENV_EXECUTION_CONCURRENCY_BURST_FACTOR, + totalConcurrencyEnabled: env.RUN_ENGINE_TOTAL_CONCURRENCY_LIMITS_ENABLED === "1", + gatesEnabled: env.RUN_ENGINE_QUEUE_GATES_ENABLED === "1", logLevel: env.RUN_ENGINE_RUN_QUEUE_LOG_LEVEL, redis: { keyPrefix: "engine:", diff --git a/apps/webapp/app/v3/runQueue.server.ts b/apps/webapp/app/v3/runQueue.server.ts index 304819f9f36..2005b059be4 100644 --- a/apps/webapp/app/v3/runQueue.server.ts +++ b/apps/webapp/app/v3/runQueue.server.ts @@ -40,6 +40,23 @@ export async function updateQueueConcurrencyLimits( await engine.runQueue.updateQueueConcurrencyLimits(environment, queueName, concurrency); } +/** Updates the RunQueue total concurrency limit for a queue (the cap across all concurrency-key values) */ +export async function updateQueueTotalConcurrencyLimits( + environment: AuthenticatedEnvironment, + queueName: string, + totalConcurrency: number +) { + await engine.runQueue.updateQueueTotalConcurrencyLimits(environment, queueName, totalConcurrency); +} + +/** Removes the RunQueue total concurrency limit for a queue */ +export async function removeQueueTotalConcurrencyLimits( + environment: AuthenticatedEnvironment, + queueName: string +) { + await engine.runQueue.removeQueueTotalConcurrencyLimits(environment, queueName); +} + /** Removes the RunQueue limits for a queue */ export async function removeQueueConcurrencyLimits( environment: AuthenticatedEnvironment, diff --git a/apps/webapp/app/v3/services/changeCurrentDeployment.server.ts b/apps/webapp/app/v3/services/changeCurrentDeployment.server.ts index f2b94ae1c05..84d406556cb 100644 --- a/apps/webapp/app/v3/services/changeCurrentDeployment.server.ts +++ b/apps/webapp/app/v3/services/changeCurrentDeployment.server.ts @@ -7,6 +7,7 @@ import { invalidateOrganizationProjectRuntimeUpdateCache } from "~/services/proj import { syncTaskIdentifiers } from "~/services/taskIdentifierRegistry.server"; import { type TaskMetadataCache, + parseTaskGates, type TaskMetadataEntry, } from "~/services/taskMetadataCache.server"; import { taskMetadataCacheInstance } from "~/services/taskMetadataCacheInstance.server"; @@ -131,6 +132,7 @@ export class ChangeCurrentDeploymentService extends BaseService { slug: true, triggerSource: true, ttl: true, + gates: true, queue: { select: { id: true, name: true } }, }, }) @@ -169,6 +171,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/concurrencyLimitNames.server.ts b/apps/webapp/app/v3/services/concurrencyLimitNames.server.ts new file mode 100644 index 00000000000..7035aaf5848 --- /dev/null +++ b/apps/webapp/app/v3/services/concurrencyLimitNames.server.ts @@ -0,0 +1,23 @@ +import { createHash } from "node:crypto"; +import { sanitizeQueueName } from "~/models/taskQueue.server"; + +/** Queue rows that back named concurrency limits live under this reserved prefix so + * they can never collide with a user's queue names. */ +export const CONCURRENCY_LIMIT_QUEUE_PREFIX = "limit/"; + +/** + * Row name for a task's anonymous inline limit. Task ids are not charset-restricted, + * so when sanitization would be lossy (or the name would overflow the 128-char queue + * name limit) a hash of the raw id keeps distinct task ids on distinct rows. Lookup + * by the public task/ name derives the row name with this same function. + */ +export function anonymousConcurrencyLimitQueueName(taskId: string): string { + const sanitized = sanitizeQueueName(taskId); + const name = `${CONCURRENCY_LIMIT_QUEUE_PREFIX}task/${sanitized}`; + if (sanitized === taskId && name.length <= 128) { + return name; + } + const hash = createHash("sha256").update(taskId).digest("hex").slice(0, 8); + const budget = 128 - `${CONCURRENCY_LIMIT_QUEUE_PREFIX}task/`.length - hash.length - 1; + return `${CONCURRENCY_LIMIT_QUEUE_PREFIX}task/${sanitized.slice(0, budget)}-${hash}`; +} diff --git a/apps/webapp/app/v3/services/concurrencyLimitsSystem.server.ts b/apps/webapp/app/v3/services/concurrencyLimitsSystem.server.ts new file mode 100644 index 00000000000..359f6f50966 --- /dev/null +++ b/apps/webapp/app/v3/services/concurrencyLimitsSystem.server.ts @@ -0,0 +1,589 @@ +import type { TaskQueue, User } from "@trigger.dev/database"; +import { errAsync, fromPromise, okAsync, type ResultAsync } from "neverthrow"; +import { Prisma, type PrismaClientOrTransaction } from "~/db.server"; +import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; +import { + removeQueueConcurrencyLimits, + removeQueueTotalConcurrencyLimits, + updateQueueConcurrencyLimits, + updateQueueTotalConcurrencyLimits, +} from "../runQueue.server"; +import { engine } from "../runEngine.server"; +import { sanitizeQueueName } from "~/models/taskQueue.server"; +import { anonymousConcurrencyLimitQueueName } from "./concurrencyLimitNames.server"; + +export type ConcurrencyLimitsSystemOptions = { + db: PrismaClientOrTransaction; + reader: PrismaClientOrTransaction; +}; + +/** Queue rows that back named concurrency limits live under this reserved prefix. */ +const LIMIT_QUEUE_PREFIX = "limit/"; + +const LIMIT_NAME_PATTERN = /^[a-zA-Z0-9_/-]{1,122}$/; + +type ConcurrencyLimitBoundValue = { + current: number | null; + base: number | null; + override: number | null; + overriddenAt: Date | null; +}; + +type ConcurrencyLimitItem = { + id: string; + name: string; + perKey: ConcurrencyLimitBoundValue; + total: ConcurrencyLimitBoundValue; + running: number; + queued: number; +}; + +/** + * An override changes only the given bounds; each bound is a non-negative integer + * (zero blocks every run holding the limit, which is how a limit is paused). + */ +type ConcurrencyLimitOverrideInput = { + perKey?: number; + total?: number; +}; + +export class ConcurrencyLimitsSystem { + constructor(private readonly options: ConcurrencyLimitsSystemOptions) {} + + private get db() { + return this.options.db; + } + + private get reader() { + return this.options.reader; + } + + get limits() { + return { + list: (environment: AuthenticatedEnvironment, page: { page: number; perPage: number }) => { + return fromPromise( + this.reader.taskQueue.findMany({ + where: limitRowsWhere(environment), + orderBy: { name: "asc" }, + skip: (page.page - 1) * page.perPage, + take: page.perPage, + }), + (error) => ({ type: "other" as const, cause: error }) + ).andThen((rows) => + fromPromise(toLimitItems(environment, rows), (error) => ({ + type: "other" as const, + cause: error, + })) + ); + }, + totalCount: (environment: AuthenticatedEnvironment) => { + return fromPromise( + this.reader.taskQueue.count({ where: limitRowsWhere(environment) }), + (error) => ({ type: "other" as const, cause: error }) + ); + }, + retrieve: (environment: AuthenticatedEnvironment, name: string) => { + return findLimitByName(this.db, environment, name).andThen((row) => + fromPromise(toLimitItems(environment, [row]), (error) => ({ + type: "other" as const, + cause: error, + })).map((items) => items[0]) + ); + }, + override: ( + environment: AuthenticatedEnvironment, + name: string, + override: ConcurrencyLimitOverrideInput, + overriddenBy?: User + ) => { + if (override.perKey === undefined && override.total === undefined) { + return errAsync({ + type: "invalid_override" as const, + message: "Provide at least one of `perKey` or `total`", + }); + } + + for (const [field, value] of Object.entries(override)) { + if (value === undefined) continue; + if (!Number.isInteger(value) || value < 0 || value > 100000) { + return errAsync({ + type: "invalid_override" as const, + message: `\`${field}\` must be an integer between 0 and 100000`, + }); + } + if (value > environment.maximumConcurrencyLimit) { + return errAsync({ + type: "invalid_override" as const, + message: `\`${field}\` (${value}) cannot exceed the environment limit (${environment.maximumConcurrencyLimit})`, + }); + } + } + + return findLimitByName(this.db, environment, name) + .andThen((row) => applyLimitOverride(this.db, row, override, overriddenBy)) + .andThen((row) => + syncLimitToEngine(environment, row) + .andThen(() => + compensateEngineFromFreshRow(this.db, environment, row.id, { + alreadySynced: { + perKey: row.concurrencyLimit, + total: row.totalConcurrencyLimit, + paused: row.paused, + }, + }) + .orElse(() => okAsync(undefined)) + .map(() => row) + ) + .orElse((error) => + compensateEngineFromFreshRow(this.db, environment, row.id) + .orElse(() => okAsync(undefined)) + .andThen(() => errAsync(error)) + ) + ) + .andThen((row) => + fromPromise(toLimitItems(environment, [row]), (error) => ({ + type: "other" as const, + cause: error, + })).map((items) => items[0]) + ); + }, + reset: (environment: AuthenticatedEnvironment, name: string) => { + return findLimitByName(this.db, environment, name) + .andThen((row) => + syncResetToEngine(environment, row).orElse((error) => + error.type === "limit_not_overridden" + ? errAsync(error) + : compensateEngineFromFreshRow(this.db, environment, row.id) + .orElse(() => okAsync(undefined)) + .andThen(() => errAsync(error)) + ) + ) + .andThen((row) => + resetLimitOverrides(this.db, row).orElse((error) => + compensateEngineFromFreshRow(this.db, environment, row.id) + .orElse(() => okAsync(undefined)) + .andThen(() => errAsync(error)) + ) + ) + .andThen((row) => + fromPromise(toLimitItems(environment, [row]), (error) => ({ + type: "other" as const, + cause: error, + })).map((items) => items[0]) + ); + }, + }; + } +} + +function concurrencyLimitDisplayId(row: Pick): string { + return `climit_${row.friendlyId.replace(/^queue_/, "")}`; +} + +function concurrencyLimitNameFromRow(row: Pick): string { + return row.name.startsWith(LIMIT_QUEUE_PREFIX) + ? row.name.slice(LIMIT_QUEUE_PREFIX.length) + : row.name; +} + +/** + * Limits live in two places: named and shared-queue inline limits are LIMIT-role + * rows under the `limit/` prefix, while an inline limit on a task's own default + * queue compiles onto that V2 QUEUE row (the design's zero-gate-cost case). Its + * derived `task/` name resolves here too, so every declared limit is + * retrievable and overridable through this one surface. V1 queue rows never + * match: their limit is queue surface, managed through the queues API. Only the + * anonymous `limit/task/` namespace requires bounds — a boundless row there is + * retired (its inline limit moved onto the task's own queue) and must fall + * through to the live queue row — while a boundless NAMED limit is a real, + * deliberately uncapped row (referenced without a declaration) that stays + * visible and cappable. + */ +function limitRowsWhere(environment: AuthenticatedEnvironment) { + return { + runtimeEnvironmentId: environment.id, + OR: [ + { + role: "LIMIT" as const, + OR: [ + { name: { not: { startsWith: `${LIMIT_QUEUE_PREFIX}task/` } } }, + { concurrencyLimit: { not: null } }, + { totalConcurrencyLimit: { not: null } }, + ], + }, + { + role: "QUEUE" as const, + concurrencyVersion: "V2" as const, + OR: [{ concurrencyLimit: { not: null } }, { totalConcurrencyLimit: { not: null } }], + }, + ], + }; +} + +function findLimitByName( + db: PrismaClientOrTransaction, + environment: AuthenticatedEnvironment, + name: string +) { + /** + * Anonymous task limits are addressed as task/, but task ids are not + * charset-restricted, so their row names are derived (sanitized, hashed when + * sanitization is lossy or the name overflows). Derive with the same function + * materialization uses so every documented name resolves; declared names keep the + * strict pattern, under which the public name IS the row suffix. + */ + const isTaskName = name.startsWith("task/"); + let limitRowName: string; + let queueRowName: string | null = null; + if (isTaskName) { + const taskId = name.slice("task/".length); + if (taskId.length === 0) { + return errAsync({ type: "limit_not_found" as const }); + } + limitRowName = anonymousConcurrencyLimitQueueName(taskId); + queueRowName = sanitizeQueueName(name); + } else { + if (!LIMIT_NAME_PATTERN.test(name)) { + return errAsync({ type: "limit_not_found" as const }); + } + limitRowName = `${LIMIT_QUEUE_PREFIX}${name}`; + } + + return fromPromise( + db.taskQueue.findFirst({ + where: { + runtimeEnvironmentId: environment.id, + name: limitRowName, + role: "LIMIT", + ...(isTaskName + ? { + OR: [{ concurrencyLimit: { not: null } }, { totalConcurrencyLimit: { not: null } }], + } + : {}), + }, + }), + (error) => ({ type: "other" as const, cause: error }) + ).andThen((row) => { + if (row) { + return okAsync(row); + } + if (queueRowName === null) { + return errAsync({ type: "limit_not_found" as const }); + } + return fromPromise( + db.taskQueue.findFirst({ + where: { + runtimeEnvironmentId: environment.id, + name: queueRowName, + role: "QUEUE", + concurrencyVersion: "V2", + }, + }), + (error) => ({ type: "other" as const, cause: error }) + ).andThen((queueRow) => { + if (!queueRow) { + return errAsync({ type: "limit_not_found" as const }); + } + return okAsync(queueRow); + }); + }); +} + +/** + * LIMIT rows read the gate machinery (group set + per-gate queued counter); a + * default-queue inline limit is its home queue, so the queue's own concurrency + * and length ARE the runs holding and waiting on the limit. + */ +async function toLimitItems( + environment: AuthenticatedEnvironment, + rows: TaskQueue[] +): Promise { + const limitNames = rows.filter((row) => row.role === "LIMIT").map((row) => row.name); + const queueNames = rows.filter((row) => row.role === "QUEUE").map((row) => row.name); + const [gateRunning, gateQueued, queueRunning, queueQueued] = await Promise.all([ + engine.totalConcurrencyOfQueues(environment, limitNames), + engine.gateQueuedCountOfQueues(environment, limitNames), + queueNames.length > 0 + ? engine.currentConcurrencyOfQueues(environment, queueNames) + : Promise.resolve({} as Record), + queueNames.length > 0 + ? engine.lengthOfQueues(environment, queueNames) + : Promise.resolve({} as Record), + ]); + const running = { ...queueRunning, ...gateRunning }; + const queued = { ...queueQueued, ...gateQueued }; + + return rows.map((row) => ({ + id: concurrencyLimitDisplayId(row), + name: concurrencyLimitNameFromRow(row), + perKey: toBound( + row.concurrencyLimit, + row.concurrencyLimitBase, + row.concurrencyLimitOverriddenAt, + environment.maximumConcurrencyLimit + ), + total: toBound( + row.totalConcurrencyLimit, + row.totalConcurrencyLimitBase, + row.totalConcurrencyLimitOverriddenAt, + environment.maximumConcurrencyLimit + ), + running: running[row.name] ?? 0, + queued: queued[row.name] ?? 0, + })); +} + +/** + * Rows store bounds exactly as declared, preserving intent when the environment limit + * later changes; the engine admits at most the environment limit regardless. `current` + * is documented as the value enforced right now, so it clamps here while `base` and + * `override` stay raw. + */ +function toBound( + stored: number | null, + base: number | null, + overriddenAt: Date | null, + environmentMaximum: number +): ConcurrencyLimitBoundValue { + const overridden = overriddenAt !== null; + return { + current: stored === null ? null : Math.min(stored, environmentMaximum), + base: overridden ? base : stored, + override: overridden ? stored : null, + overriddenAt, + }; +} + +function applyLimitOverride( + db: PrismaClientOrTransaction, + row: TaskQueue, + override: ConcurrencyLimitOverrideInput, + overriddenBy?: User +) { + const now = new Date(); + const data: Record = {}; + + if (override.perKey !== undefined) { + data.concurrencyLimit = override.perKey; + data.concurrencyLimitBase = row.concurrencyLimitOverriddenAt + ? row.concurrencyLimitBase + : (row.concurrencyLimit ?? null); + data.concurrencyLimitOverriddenAt = now; + data.concurrencyLimitOverriddenBy = overriddenBy?.id ?? null; + data.concurrencyLimitOverridePercent = null; + } + + if (override.total !== undefined) { + data.totalConcurrencyLimit = override.total; + data.totalConcurrencyLimitBase = row.totalConcurrencyLimitOverriddenAt + ? row.totalConcurrencyLimitBase + : (row.totalConcurrencyLimit ?? null); + data.totalConcurrencyLimitOverriddenAt = now; + data.totalConcurrencyLimitOverriddenBy = overriddenBy?.id ?? null; + } + + return guardedLimitUpdate(db, row, data); +} + +/** + * Enforce first, then persist: the engine syncs to the declared base BEFORE the + * override markers clear, so an engine failure leaves the markers set and a retry + * converges instead of being rejected while the overridden limit stays enforced. + */ +/** + * A paused row's pause IS the engine per-key value 0 (the DB concurrencyLimit + * column keeps the configured value), so every per-key engine write from this + * surface must preserve it — otherwise an override or reset that only touched + * `total` would silently resume a queue every other surface still reports as + * paused. Named LIMIT rows are never paused (pausing a limit is an override to + * `{ total: 0 }`), so they always take the target branch. + */ +function perKeyEngineWrite( + environment: AuthenticatedEnvironment, + row: Pick, + target: number | null | undefined +) { + if (row.paused) { + return updateQueueConcurrencyLimits(environment, row.name, 0); + } + return typeof target === "number" + ? updateQueueConcurrencyLimits(environment, row.name, target) + : removeQueueConcurrencyLimits(environment, row.name); +} + +function syncResetToEngine( + environment: AuthenticatedEnvironment, + row: TaskQueue +): ResultAsync< + TaskQueue, + { type: "limit_not_overridden" } | { type: "sync_limit_to_engine_failed"; cause: unknown } +> { + if (row.concurrencyLimitOverriddenAt === null && row.totalConcurrencyLimitOverriddenAt === null) { + return errAsync({ type: "limit_not_overridden" as const }); + } + + const perKeyTarget = row.concurrencyLimitOverriddenAt + ? row.concurrencyLimitBase + : row.concurrencyLimit; + const totalTarget = row.totalConcurrencyLimitOverriddenAt + ? row.totalConcurrencyLimitBase + : row.totalConcurrencyLimit; + + const perKeySync = perKeyEngineWrite(environment, row, perKeyTarget); + + const totalSync = + typeof totalTarget === "number" + ? updateQueueTotalConcurrencyLimits(environment, row.name, totalTarget) + : removeQueueTotalConcurrencyLimits(environment, row.name); + + return fromPromise(settleBothEngineWrites(perKeySync, totalSync), (error) => ({ + type: "sync_limit_to_engine_failed" as const, + cause: error, + })).map(() => row); +} + +function resetLimitOverrides(db: PrismaClientOrTransaction, row: TaskQueue) { + const data: Record = {}; + + if (row.concurrencyLimitOverriddenAt !== null) { + data.concurrencyLimit = row.concurrencyLimitBase; + data.concurrencyLimitBase = null; + data.concurrencyLimitOverriddenAt = null; + data.concurrencyLimitOverriddenBy = null; + data.concurrencyLimitOverridePercent = null; + } + + if (row.totalConcurrencyLimitOverriddenAt !== null) { + data.totalConcurrencyLimit = row.totalConcurrencyLimitBase; + data.totalConcurrencyLimitBase = null; + data.totalConcurrencyLimitOverriddenAt = null; + data.totalConcurrencyLimitOverriddenBy = null; + } + + return guardedLimitUpdate(db, row, data); +} + +/** + * Both engine writes settle before a failure is reported, so no write is still in + * flight when a caller's compensation runs — a late sibling can never land after + * the compensating re-sync and leave one bound stale. + */ +async function settleBothEngineWrites(a: Promise, b: Promise): Promise { + const results = await Promise.allSettled([a, b]); + const failed = results.find((result) => result.status === "rejected"); + if (failed && failed.status === "rejected") { + throw failed.reason; + } +} + +/** + * Optimistic update: the where clause carries the row's updatedAt plus both override + * markers as read, so ANY concurrent write — another override or reset, or a deploy + * refreshing the declared values — makes this update miss (P2025) and the caller + * gets a conflict instead of persisting values computed from a stale row. The + * markers narrow the same-millisecond updatedAt window to writes that also leave + * both markers untouched. + */ +function guardedLimitUpdate( + db: PrismaClientOrTransaction, + row: TaskQueue, + data: Record +) { + return fromPromise( + db.taskQueue.update({ + where: { + id: row.id, + updatedAt: row.updatedAt, + concurrencyLimitOverriddenAt: row.concurrencyLimitOverriddenAt, + totalConcurrencyLimitOverriddenAt: row.totalConcurrencyLimitOverriddenAt, + }, + data, + }), + (error) => { + if (error instanceof Prisma.PrismaClientKnownRequestError && error.code === "P2025") { + return { type: "conflict" as const }; + } + return { type: "limit_update_failed" as const, cause: error }; + } + ); +} + +type SyncedLimitValues = { perKey: number | null; total: number | null; paused: boolean }; + +/** + * Re-syncs the engine from fresh reads of the row until the enforced values stop + * moving (bounded), the same convergence the deploy sync uses: every actor writes + * Postgres before its own engine sync, so re-syncing whatever is freshest + * converges. The fixpoint compares the values the engine enforces rather than + * updatedAt, because Prisma's @updatedAt has millisecond precision and two writes + * in the same millisecond are indistinguishable by timestamp. Matching values + * only prove this actor once synced them, not that the engine still holds them + * (another actor may have diverged it and written the same values back), but + * skipping keeps divergence non-silent: a failed engine write always surfaces + * to that actor's caller, which can retry, and a stale write landing after the + * loop's final read (the loop is bounded) is healed by the next sync or deploy, + * the same residual the deploy-time queue sync accepts. Callers use it two ways: after a failure (a reset's enforce-first engine write preceding a + * persist that then conflicts, or an override's sync failing after its persist), + * where the original error still reaches the caller; and after a successful sync + * with `alreadySynced` set to the values just synced, where an unchanged row + * costs one read and a moved row is re-synced. + */ +function compensateEngineFromFreshRow( + db: PrismaClientOrTransaction, + environment: AuthenticatedEnvironment, + rowId: string, + options?: { alreadySynced?: SyncedLimitValues } +) { + return fromPromise( + (async () => { + let lastSynced: SyncedLimitValues | null = options?.alreadySynced ?? null; + for (let i = 0; i < 3; i++) { + const fresh = await db.taskQueue.findFirst({ where: { id: rowId } }); + if ( + !fresh || + (lastSynced !== null && + fresh.concurrencyLimit === lastSynced.perKey && + fresh.totalConcurrencyLimit === lastSynced.total && + fresh.paused === lastSynced.paused) + ) { + return; + } + await settleBothEngineWrites( + perKeyEngineWrite(environment, fresh, fresh.concurrencyLimit), + typeof fresh.totalConcurrencyLimit === "number" + ? updateQueueTotalConcurrencyLimits( + environment, + fresh.name, + fresh.totalConcurrencyLimit + ) + : removeQueueTotalConcurrencyLimits(environment, fresh.name) + ); + lastSynced = { + perKey: fresh.concurrencyLimit, + total: fresh.totalConcurrencyLimit, + paused: fresh.paused, + }; + } + })(), + (error) => ({ type: "other" as const, cause: error }) + ); +} + +/** + * Pushes both engine keys from the row: the per-key limit and the total. Limit + * rows are never paused (pausing a limit is an override to `{ total: 0 }`), so + * both keys sync unconditionally, unlike queue rows. + */ +function syncLimitToEngine(environment: AuthenticatedEnvironment, row: TaskQueue) { + const perKeySync = perKeyEngineWrite(environment, row, row.concurrencyLimit); + + const totalSync = + typeof row.totalConcurrencyLimit === "number" + ? updateQueueTotalConcurrencyLimits(environment, row.name, row.totalConcurrencyLimit) + : removeQueueTotalConcurrencyLimits(environment, row.name); + + return fromPromise(settleBothEngineWrites(perKeySync, totalSync), (error) => ({ + type: "sync_limit_to_engine_failed" as const, + cause: error, + })).map(() => row); +} diff --git a/apps/webapp/app/v3/services/concurrencyLimitsSystemInstance.server.ts b/apps/webapp/app/v3/services/concurrencyLimitsSystemInstance.server.ts new file mode 100644 index 00000000000..eeaa6882e9f --- /dev/null +++ b/apps/webapp/app/v3/services/concurrencyLimitsSystemInstance.server.ts @@ -0,0 +1,15 @@ +import { prisma, $replica } from "~/db.server"; +import { ConcurrencyLimitsSystem } from "./concurrencyLimitsSystem.server"; +import { singleton } from "~/utils/singleton"; + +export const concurrencyLimitsSystem = singleton( + "concurrency-limits-system", + initializeConcurrencyLimitsSystemInstance +); + +function initializeConcurrencyLimitsSystemInstance() { + return new ConcurrencyLimitsSystem({ + db: prisma, + reader: $replica, + }); +} diff --git a/apps/webapp/app/v3/services/concurrencySystem.server.ts b/apps/webapp/app/v3/services/concurrencySystem.server.ts index 51c51674234..b3dfff8176c 100644 --- a/apps/webapp/app/v3/services/concurrencySystem.server.ts +++ b/apps/webapp/app/v3/services/concurrencySystem.server.ts @@ -1,9 +1,14 @@ import type { TaskQueue, User } from "@trigger.dev/database"; -import { errAsync, fromPromise, okAsync } from "neverthrow"; -import type { PrismaClientOrTransaction } from "~/db.server"; +import { errAsync, fromPromise, okAsync, type ResultAsync } from "neverthrow"; +import { Prisma, type PrismaClientOrTransaction } from "~/db.server"; import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; import { logger } from "~/services/logger.server"; -import { removeQueueConcurrencyLimits, updateQueueConcurrencyLimits } from "../runQueue.server"; +import { + removeQueueConcurrencyLimits, + removeQueueTotalConcurrencyLimits, + updateQueueConcurrencyLimits, + updateQueueTotalConcurrencyLimits, +} from "../runQueue.server"; import { engine } from "../runEngine.server"; export type ConcurrencySystemOptions = { @@ -62,21 +67,62 @@ export class ConcurrencySystem { environment: AuthenticatedEnvironment, queue: QueueInput, override: ConcurrencyLimitOverride, - overriddenBy?: User + overriddenBy?: User, + opts?: QueueMutationOpts ) => { return findQueueFromInput(this.db, environment, queue) + .andThen((queue) => guardQueueVersion(queue, opts)) .andThen((queue) => overrideQueueConcurrencyLimit(this.db, environment, queue, override, overriddenBy) ) .andThen((queue) => syncQueueConcurrencyToEngine(environment, queue)) .andThen((queue) => getQueueStats(environment, queue)); }, - resetConcurrencyLimit: (environment: AuthenticatedEnvironment, queue: QueueInput) => { + resetConcurrencyLimit: ( + environment: AuthenticatedEnvironment, + queue: QueueInput, + opts?: QueueMutationOpts + ) => { return findQueueFromInput(this.db, environment, queue) + .andThen((queue) => guardQueueVersion(queue, opts)) .andThen((queue) => resetQueueConcurrencyLimit(this.db, queue)) .andThen((queue) => syncQueueConcurrencyToEngine(environment, queue)) .andThen((queue) => getQueueStats(environment, queue)); }, + overrideTotalConcurrencyLimit: ( + environment: AuthenticatedEnvironment, + queue: QueueInput, + totalConcurrencyLimit: number, + overriddenBy?: User + ) => { + return findQueueFromInput(this.db, environment, queue) + .andThen((queue) => + overrideQueueTotalConcurrencyLimit( + this.db, + environment, + queue, + totalConcurrencyLimit, + overriddenBy + ) + ) + .andThen((queue) => syncQueueTotalConcurrencyToEngine(environment, queue)) + .andThen((queue) => getQueueStats(environment, queue)); + }, + resetTotalConcurrencyLimit: (environment: AuthenticatedEnvironment, queue: QueueInput) => { + return findQueueFromInput(this.db, environment, queue) + .andThen((queue) => syncQueueTotalConcurrencyResetToEngine(environment, queue)) + .andThen((queue) => + resetQueueTotalConcurrencyLimit(this.db, queue).orElse((error) => + error.type === "concurrent_modification" + ? healTotalConcurrencyFromRow(this.db, environment, queue.id).andThen(() => + errAsync(error) + ) + : errAsync(error) + ) + ) + .andThen((queue) => syncQueueTotalConcurrencyToEngine(environment, queue)) + .andThen((queue) => getQueueStats(environment, queue)); + }, /** * Recalculates the materialized limit of every percent-based override in the environment * against its CURRENT maximumConcurrencyLimit and syncs changed queues to the run engine. @@ -149,6 +195,22 @@ function findQueueFromInput( return findQueueByName(db, environment, queueName); } +/** + * The public queue override/reset endpoints are the V1 lever; a V2 queue's + * limits are managed through the concurrency-limits endpoints (a default-queue + * inline limit under its derived task/ name), and a V2 response hides + * queue-level concurrency, so mutating one here would succeed invisibly. The + * dashboard's own actions pass no opts and keep working on every version. + */ +type QueueMutationOpts = { v1Only?: boolean }; + +function guardQueueVersion(queue: TaskQueue, opts: QueueMutationOpts | undefined) { + if (opts?.v1Only && queue.concurrencyVersion === "V2") { + return errAsync({ type: "queue_version_unsupported" as const }); + } + return okAsync(queue); +} + function findQueueByFriendlyId( db: PrismaClientOrTransaction, environment: AuthenticatedEnvironment, @@ -159,6 +221,7 @@ function findQueueByFriendlyId( where: { runtimeEnvironmentId: environment.id, friendlyId, + role: "QUEUE", }, }), (error) => ({ @@ -183,6 +246,7 @@ function findQueueByName( where: { runtimeEnvironmentId: environment.id, name: queue, + role: "QUEUE", }, }), (error) => ({ @@ -197,6 +261,22 @@ function findQueueByName( }); } +/** + * Maps a mutation's update failure. P2025 means the optimistic marker in the where + * clause moved between this mutation's read and its write (a concurrent override or + * reset won); surfacing it keeps the loser from persisting values derived from the + * stale read, which would silently corrupt the saved base limit. + */ +function queueUpdateError(error: unknown) { + if (error instanceof Prisma.PrismaClientKnownRequestError && error.code === "P2025") { + return { + type: "concurrent_modification" as const, + message: "The queue's concurrency was changed by another request; retry.", + }; + } + return { type: "queue_update_failed" as const, cause: error }; +} + function overrideQueueConcurrencyLimit( db: PrismaClientOrTransaction, environment: AuthenticatedEnvironment, @@ -253,6 +333,7 @@ function overrideQueueConcurrencyLimit( db.taskQueue.update({ where: { id: queue.id, + concurrencyLimitOverriddenAt: queue.concurrencyLimitOverriddenAt, }, data: { concurrencyLimit: newConcurrencyLimit, @@ -262,10 +343,7 @@ function overrideQueueConcurrencyLimit( concurrencyLimitOverriddenBy: overriddenBy?.id ?? null, }, }), - (error) => ({ - type: "queue_update_failed" as const, - cause: error, - }) + queueUpdateError ); } @@ -278,7 +356,7 @@ function resetQueueConcurrencyLimit(db: PrismaClientOrTransaction, queue: TaskQu return fromPromise( db.taskQueue.update({ - where: { id: queue.id }, + where: { id: queue.id, concurrencyLimitOverriddenAt: queue.concurrencyLimitOverriddenAt }, data: { concurrencyLimitOverriddenAt: null, concurrencyLimit: newConcurrencyLimit, @@ -287,10 +365,7 @@ function resetQueueConcurrencyLimit(db: PrismaClientOrTransaction, queue: TaskQu concurrencyLimitOverriddenBy: null, }, }), - (error) => ({ - type: "queue_update_failed" as const, - cause: error, - }) + queueUpdateError ); } @@ -316,6 +391,181 @@ function syncQueueConcurrencyToEngine(environment: AuthenticatedEnvironment, que } } +function overrideQueueTotalConcurrencyLimit( + db: PrismaClientOrTransaction, + environment: AuthenticatedEnvironment, + queue: TaskQueue, + totalConcurrencyLimit: number, + overriddenBy?: User +) { + const maximum = environment.maximumConcurrencyLimit; + + if (!Number.isFinite(totalConcurrencyLimit) || totalConcurrencyLimit < 0) { + return errAsync({ + type: "invalid_override" as const, + message: "Combined concurrency limit must be a non-negative number", + }); + } + + if (totalConcurrencyLimit > maximum) { + return errAsync({ + type: "concurrency_limit_exceeds_maximum" as const, + message: `Combined concurrency limit (${totalConcurrencyLimit}) cannot exceed the environment limit (${maximum})`, + }); + } + + const totalConcurrencyLimitBase = queue.totalConcurrencyLimitOverriddenAt + ? queue.totalConcurrencyLimitBase + : queue.totalConcurrencyLimit; + + return fromPromise( + db.taskQueue.update({ + where: { + id: queue.id, + totalConcurrencyLimitOverriddenAt: queue.totalConcurrencyLimitOverriddenAt, + }, + data: { + totalConcurrencyLimit, + totalConcurrencyLimitBase: totalConcurrencyLimitBase ?? null, + totalConcurrencyLimitOverriddenAt: new Date(), + totalConcurrencyLimitOverriddenBy: overriddenBy?.id ?? null, + }, + }), + queueUpdateError + ); +} + +/** + * Enforce first, then persist: syncs the engine to the declared base BEFORE clearing + * the override marker, so an engine failure leaves the marker set and a retry + * converges instead of being rejected while the overridden limit stays enforced. + */ +function syncQueueTotalConcurrencyResetToEngine( + environment: AuthenticatedEnvironment, + queue: TaskQueue +) { + if (queue.totalConcurrencyLimitOverriddenAt === null) { + return errAsync({ type: "queue_not_overridden" as const }); + } + + if (typeof queue.totalConcurrencyLimitBase === "number") { + return fromPromise( + updateQueueTotalConcurrencyLimits(environment, queue.name, queue.totalConcurrencyLimitBase), + (error) => ({ + type: "sync_queue_concurrency_to_engine_failed" as const, + cause: error, + }) + ).andThen(() => okAsync(queue)); + } + + return fromPromise(removeQueueTotalConcurrencyLimits(environment, queue.name), (error) => ({ + type: "sync_queue_concurrency_to_engine_failed" as const, + cause: error, + })).andThen(() => okAsync(queue)); +} + +/** + * A guarded reset conflict can leave the engine holding the reset target written by the + * enforce-first sync while the DB retains the concurrent winner's override (whose own + * final sync may itself have failed). Re-sync the engine from the fresh row so the + * enforced value tracks the persisted one; best effort, the conflict error is returned + * either way. + */ +function healTotalConcurrencyFromRow( + db: PrismaClientOrTransaction, + environment: AuthenticatedEnvironment, + queueId: string +) { + return fromPromise( + (async () => { + /** + * Bounded convergence, mirroring compensateEngineFromFreshRow: re-read after each + * engine write and stop once the persisted value held still, so a mutation that + * commits and syncs between this heal's read and its write is re-applied instead + * of being rolled back by the heal's stale write. A stale write landing after the + * loop's final read remains possible and is healed by the next sync or deploy. + */ + let lastSynced: number | null | undefined; + for (let i = 0; i < 3; i++) { + const fresh = await db.taskQueue.findFirst({ where: { id: queueId } }); + if (!fresh || (lastSynced !== undefined && fresh.totalConcurrencyLimit === lastSynced)) { + return; + } + if (typeof fresh.totalConcurrencyLimit === "number") { + await updateQueueTotalConcurrencyLimits( + environment, + fresh.name, + fresh.totalConcurrencyLimit + ); + } else { + await removeQueueTotalConcurrencyLimits(environment, fresh.name); + } + lastSynced = fresh.totalConcurrencyLimit; + } + })().catch((error) => { + logger.error("Failed to re-sync total concurrency after a reset conflict", { + error, + queueId, + }); + }), + () => ({ type: "other" as const }) + ).orElse(() => okAsync(undefined)); +} + +function resetQueueTotalConcurrencyLimit( + db: PrismaClientOrTransaction, + queue: TaskQueue +): ResultAsync< + TaskQueue, + | { type: "queue_not_overridden" } + | { type: "concurrent_modification"; message: string } + | { type: "queue_update_failed"; cause: unknown } +> { + if (queue.totalConcurrencyLimitOverriddenAt === null) { + return errAsync({ type: "queue_not_overridden" as const }); + } + + return fromPromise( + db.taskQueue.update({ + where: { + id: queue.id, + totalConcurrencyLimitOverriddenAt: queue.totalConcurrencyLimitOverriddenAt, + }, + data: { + totalConcurrencyLimit: queue.totalConcurrencyLimitBase, + totalConcurrencyLimitBase: null, + totalConcurrencyLimitOverriddenAt: null, + totalConcurrencyLimitOverriddenBy: null, + }, + }), + queueUpdateError + ); +} + +/** + * The total limit key is separate from the per-queue limit key that pause zeroes, + * so it syncs regardless of the paused state. + */ +function syncQueueTotalConcurrencyToEngine( + environment: AuthenticatedEnvironment, + queue: TaskQueue +) { + if (typeof queue.totalConcurrencyLimit === "number") { + return fromPromise( + updateQueueTotalConcurrencyLimits(environment, queue.name, queue.totalConcurrencyLimit), + (error) => ({ + type: "sync_queue_concurrency_to_engine_failed" as const, + cause: error, + }) + ).andThen(() => okAsync(queue)); + } + + return fromPromise(removeQueueTotalConcurrencyLimits(environment, queue.name), (error) => ({ + type: "sync_queue_concurrency_to_engine_failed" as const, + cause: error, + })).andThen(() => okAsync(queue)); +} + function getQueueStats(environment: AuthenticatedEnvironment, queue: TaskQueue) { return fromPromise( Promise.all([ diff --git a/apps/webapp/app/v3/services/createBackgroundWorker.server.test.ts b/apps/webapp/app/v3/services/createBackgroundWorker.server.test.ts index ca6113aa3f0..675727db80b 100644 --- a/apps/webapp/app/v3/services/createBackgroundWorker.server.test.ts +++ b/apps/webapp/app/v3/services/createBackgroundWorker.server.test.ts @@ -214,6 +214,7 @@ describe("worker task creation", () => { triggerSource: "STANDARD", queueId: queue.id, queueName: queue.name, + gates: null, }, ], [ @@ -223,6 +224,7 @@ describe("worker task creation", () => { triggerSource: "STANDARD", queueId: queue.id, queueName: queue.name, + gates: null, }, ], ]); @@ -244,6 +246,61 @@ describe("worker task creation", () => { expect(afterRetry.description).toBe(created.description); } ); + + containerTest( + "a named limit colliding with a legacy queue in the limit/ namespace fails the deploy instead of repurposing the row", + async ({ prisma }) => { + const { project, devEnv } = await seedProjectWithEnvs(prisma); + const worker = await prisma.backgroundWorker.create({ + data: { + friendlyId: `worker_${devEnv.id}`, + contentHash: "limit-collision-content", + version: "20260811.1", + metadata: {}, + projectId: project.id, + runtimeEnvironmentId: devEnv.id, + }, + }); + /** Slashes were historically legal in user queue names, so a QUEUE-role row can + * already occupy the reserved name a declared limit would materialize under. */ + const legacyQueue = await prisma.taskQueue.create({ + data: { + friendlyId: `queue_${devEnv.id}_legacy`, + name: "limit/openai", + type: "NAMED", + version: "V2", + concurrencyLimit: 7, + projectId: project.id, + runtimeEnvironmentId: devEnv.id, + }, + }); + const environment = { ...asEnv(devEnv), project } as AuthenticatedEnvironment; + + const metadata = { + contentHash: "limit-collision-content", + tasks: [ + { + id: "collision-task", + filePath: "src/trigger/collision-task.ts", + exportName: "collisionTask", + concurrency: { limits: ["openai"] }, + }, + ], + concurrencyLimits: [{ name: "openai", total: 5 }], + } as unknown as BackgroundWorkerMetadata; + + await expect( + prisma.$transaction((tx) => createWorkerResources(metadata, worker, environment, tx)) + ).rejects.toThrow(/collides with an existing queue/); + + const untouched = await prisma.taskQueue.findFirstOrThrow({ + where: { id: legacyQueue.id }, + }); + expect(untouched.role).toBe("QUEUE"); + expect(untouched.concurrencyLimit).toBe(7); + expect(untouched.totalConcurrencyLimit).toBeNull(); + } + ); }); describe("declarative schedule preflight", () => { diff --git a/apps/webapp/app/v3/services/createBackgroundWorker.server.ts b/apps/webapp/app/v3/services/createBackgroundWorker.server.ts index 06b2e3454fc..13a24576042 100644 --- a/apps/webapp/app/v3/services/createBackgroundWorker.server.ts +++ b/apps/webapp/app/v3/services/createBackgroundWorker.server.ts @@ -17,7 +17,13 @@ import { stringifyDuration, } from "@trigger.dev/core/v3/isomorphic"; import { randomBytes } from "node:crypto"; -import type { BackgroundWorker, TaskQueue, TaskQueueType } from "@trigger.dev/database"; +import type { + BackgroundWorker, + TaskQueue, + TaskQueueConcurrencyVersion, + TaskQueueRole, + TaskQueueType, +} from "@trigger.dev/database"; import cronstrue from "cronstrue"; import type { PrismaClientOrTransaction, WebhookDatabase } from "~/db.server"; import { $transaction, Prisma, boundedIn, webhookPrisma } from "~/db.server"; @@ -32,11 +38,17 @@ import { } from "~/services/taskMetadataCache.server"; import { taskMetadataCacheInstance } from "~/services/taskMetadataCacheInstance.server"; import { generateFriendlyId } from "../friendlyIdentifiers"; +import { + anonymousConcurrencyLimitQueueName, + CONCURRENCY_LIMIT_QUEUE_PREFIX, +} from "./concurrencyLimitNames.server"; import { engine } from "../runEngine.server"; import { removeQueueConcurrencyLimits, + removeQueueTotalConcurrencyLimits, updateEnvConcurrencyLimits, updateQueueConcurrencyLimits, + updateQueueTotalConcurrencyLimits, } from "../runQueue.server"; import { resolveScheduleWindow } from "@internal/schedule-engine"; import { scheduleEngine } from "../scheduleEngine.server"; @@ -105,6 +117,8 @@ export class CreateBackgroundWorkerService extends BaseService { }, }); + validateWorkerConcurrencyDeclarations(body.metadata); + const latestBackgroundWorker = project.backgroundWorkers[0]; // Validate before persisting a worker or accepting a matching hash from an earlier attempt. @@ -365,6 +379,7 @@ export async function createWorkerResources( // Create the queues const queues = await createWorkerQueues(metadata, worker, environment, prisma); + await createWorkerConcurrencyLimits(metadata, worker, environment, prisma); // Create the tasks const taskEntries = await createWorkerTasks( @@ -406,9 +421,139 @@ async function createWorkerTasks( if (entry) entries.push(entry); } } + + await retireStaleAnonymousConcurrencyLimitRows(metadata.tasks, environment, prisma); + return entries; } +/** + * A task whose inline limit no longer compiles to its anonymous LIMIT row (the + * limit moved onto the task's own default queue, or was removed) leaves that + * row behind, where it would shadow the live task/ name on the + * concurrency-limits surface and keep stale engine keys. Retiring clears the + * row's bounds and override state — dropping it from listing and name + * resolution so the queue-row fallback resolves — and removes its engine keys. + */ +/** + * Best-effort engine sync of a limit row's bounds (pause-aware): both writes + * settle, and any rejection is surfaced in the logs rather than swallowed — a + * silent failure here is a limit the dashboard reports but the engine does not + * enforce. + */ +async function syncLimitRowEngineStateBestEffort( + environment: AuthenticatedEnvironment, + row: Pick +): Promise { + const results = await Promise.allSettled([ + row.paused + ? updateQueueConcurrencyLimits(environment, row.name, 0) + : typeof row.concurrencyLimit === "number" + ? updateQueueConcurrencyLimits(environment, row.name, row.concurrencyLimit) + : removeQueueConcurrencyLimits(environment, row.name), + typeof row.totalConcurrencyLimit === "number" + ? updateQueueTotalConcurrencyLimits(environment, row.name, row.totalConcurrencyLimit) + : removeQueueTotalConcurrencyLimits(environment, row.name), + ]); + const failures = results.filter((result) => result.status === "rejected"); + if (failures.length > 0) { + logger.error("syncLimitRowEngineStateBestEffort: engine sync failed", { + environmentId: environment.id, + queueName: row.name, + errors: failures.map((failure) => String((failure as PromiseRejectedResult).reason)), + }); + } +} + +async function retireStaleAnonymousConcurrencyLimitRows( + tasks: TaskResource[], + environment: AuthenticatedEnvironment, + prisma: PrismaClientOrTransaction +): Promise { + const candidateNames = tasks + .filter((task) => !(task.concurrency?.inline && task.queue?.name)) + .map((task) => anonymousConcurrencyLimitQueueName(task.id)); + if (candidateNames.length === 0) { + return; + } + + const staleRows = await prisma.taskQueue.findMany({ + where: { + runtimeEnvironmentId: environment.id, + role: "LIMIT", + name: { in: boundedIn(candidateNames) }, + OR: [{ concurrencyLimit: { not: null } }, { totalConcurrencyLimit: { not: null } }], + }, + select: { id: true, name: true, updatedAt: true }, + }); + if (staleRows.length === 0) { + return; + } + + /** Engine keys go first, and the row's bounds are only nulled once both + * removals succeeded: a failed removal leaves the row bounded, so the next + * deploy retries the retirement instead of stranding a stale engine limit + * (worst case a pause-by-zero) that nothing can see or clear. Each null is + * guarded on the row's read updatedAt, so a concurrent writer (an operator + * override, or another deploy re-creating the limit) wins and its own engine + * sync governs; a racing limits-surface sync converges via its freshness + * re-check against the nulled row. */ + for (const row of staleRows) { + /** Both removals settle before anything else runs, so no removal is still + * in flight when the heal writes; on any failure the heal re-reads the row + * fresh (a concurrent writer may have moved it since the findMany) and + * re-syncs the engine to whatever is persisted now. */ + const removals = await Promise.allSettled([ + removeQueueConcurrencyLimits(environment, row.name), + removeQueueTotalConcurrencyLimits(environment, row.name), + ]); + const removalFailures = removals.filter((result) => result.status === "rejected"); + if (removalFailures.length > 0) { + logger.error( + "retireStaleAnonymousConcurrencyLimitRows: engine cleanup failed, healing and retrying next deploy", + { + environmentId: environment.id, + queueName: row.name, + errors: removalFailures.map((failure) => + String((failure as PromiseRejectedResult).reason) + ), + } + ); + const fresh = await prisma.taskQueue.findFirst({ where: { id: row.id } }); + if (fresh) { + await syncLimitRowEngineStateBestEffort(environment, fresh); + } + continue; + } + const retired = await prisma.taskQueue.updateMany({ + where: { id: row.id, updatedAt: row.updatedAt }, + data: { + concurrencyLimit: null, + concurrencyLimitBase: null, + concurrencyLimitOverriddenAt: null, + concurrencyLimitOverriddenBy: null, + concurrencyLimitOverridePercent: null, + totalConcurrencyLimit: null, + totalConcurrencyLimitBase: null, + totalConcurrencyLimitOverriddenAt: null, + totalConcurrencyLimitOverriddenBy: null, + }, + }); + if (retired.count === 0) { + /** A concurrent writer (an operator override, or another deploy) took the + * row between the read and the null, and the key removal above may have + * erased the engine state that writer just synced — including a + * pause-by-zero. Restore the engine from the fresh row so the winner's + * bounds stay enforced; both writers write the same fresh values, so the + * race converges. */ + const fresh = await prisma.taskQueue.findFirst({ where: { id: row.id } }); + if (fresh) { + await syncLimitRowEngineStateBestEffort(environment, fresh); + } + } + } +} + async function createWorkerTask( task: TaskResource, queues: Array, @@ -418,10 +563,59 @@ async function createWorkerTask( tasksToBackgroundFiles?: Map ): Promise { let queue: TaskQueue | undefined; + let compiledGates: Array<{ queue: string; concurrencyKey?: string }> = []; let resolvedTriggerSource: "SCHEDULED" | "AGENT" | "WEBHOOK" | "STANDARD" | undefined; let resolvedTtl: string | null | undefined; try { + const concurrency = task.concurrency; + + if (task.queue?.name) { + assertNotReservedQueueName(task.queue.name, `Task "${task.id}"`); + } + + if (concurrency && typeof task.queue?.concurrencyLimit === "number") { + throw new ServiceValidationError( + `Task "${task.id}" declares both a queue concurrencyLimit and the concurrency option; use concurrency.` + ); + } + + compiledGates = (concurrency?.limits ?? []).map((name) => ({ + queue: concurrencyLimitQueueName(name), + })); + + let queueConcurrencyLimit = task.queue?.concurrencyLimit; + let queueTotalConcurrencyLimit = task.queue?.combinedConcurrencyLimit; + + if (concurrency?.inline) { + if (!task.queue?.name) { + queueConcurrencyLimit = concurrency.inline.perKey; + queueTotalConcurrencyLimit = concurrency.inline.total; + } else { + if (compiledGates.length > 1) { + throw new ServiceValidationError( + `Task "${task.id}": an inline limit on a shared queue uses a gate slot, so at most one named limit can be combined with it.` + ); + } + const anonymousQueueName = anonymousConcurrencyLimitQueueName(task.id); + await createWorkerQueue( + { + name: anonymousQueueName, + concurrencyLimit: concurrency.inline.perKey ?? null, + combinedConcurrencyLimit: concurrency.inline.total ?? null, + }, + `task/${task.id}`, + "NAMED", + worker, + environment, + prisma, + "LIMIT", + "V2" + ); + compiledGates = [{ queue: anonymousQueueName }, ...compiledGates]; + } + } + queue = queues.find((queue) => queue.name === task.queue?.name); if (!queue) { @@ -429,13 +623,22 @@ async function createWorkerTask( queue = await createWorkerQueue( { name: task.queue?.name ?? `task/${task.id}`, - concurrencyLimit: task.queue?.concurrencyLimit, + concurrencyLimit: queueConcurrencyLimit, + combinedConcurrencyLimit: queueTotalConcurrencyLimit, }, - task.id, + task.queue?.name ?? task.id, task.queue?.name ? "NAMED" : "VIRTUAL", worker, environment, - prisma + prisma, + "QUEUE", + /** + * V2 marks rows whose limit fields hold the new perKey/total vocabulary. That + * only happens when an inline limit compiles into the task's own default queue; + * a shared queue keeps V1 regardless of task concurrency (which lives in gates + * and LIMIT rows), matching rows materialized from queue() declarations. + */ + concurrency?.inline && !task.queue?.name ? "V2" : "V1" ); } @@ -465,6 +668,7 @@ async function createWorkerTask( exportName: task.exportName, retryConfig: task.retry, queueConfig: task.queue, + gates: compiledGates.length > 0 ? compiledGates : task.gates, machineConfig: task.machine, triggerSource: resolvedTriggerSource, config: task.agentConfig ? (task.agentConfig as any) : undefined, @@ -503,8 +707,12 @@ async function createWorkerTask( triggerSource: resolvedTriggerSource, queueId: queue.id, queueName: queue.name, + gates: compiledGates.length > 0 ? compiledGates : (task.gates ?? null), }; } catch (error) { + if (error instanceof ServiceValidationError) { + throw error; + } if (error instanceof Prisma.PrismaClientKnownRequestError) { logger.error("Prisma Error creating background worker task", { error: { @@ -550,6 +758,7 @@ async function createWorkerQueues( const chunk = metadata.queues.slice(i, i + CHUNK_SIZE); const queueChunk = await Promise.all( chunk.map(async (queue) => { + assertNotReservedQueueName(queue.name, `Queue "${queue.name}"`); return createWorkerQueue(queue, queue.name, "NAMED", worker, environment, prisma); }) ); @@ -559,13 +768,130 @@ async function createWorkerQueues( return allQueues; } +/** + * Rejects invalid concurrency declarations before any worker rows are written, so a + * failed deploy leaves nothing behind for a same-content retry to return. + */ +export function validateWorkerConcurrencyDeclarations(metadata: BackgroundWorkerMetadata): void { + for (const queue of metadata.queues ?? []) { + assertNotReservedQueueName(queue.name, `Queue "${queue.name}"`); + } + + for (const limit of metadata.concurrencyLimits ?? []) { + assertValidConcurrencyLimitName(limit.name); + } + + for (const task of metadata.tasks) { + if (task.queue?.name) { + assertNotReservedQueueName(task.queue.name, `Task "${task.id}"`); + } + + for (const gate of task.gates ?? []) { + assertNotReservedQueueName(gate.queue, `Task "${task.id}" gate "${gate.queue}"`); + } + + const concurrency = task.concurrency; + if (!concurrency) { + continue; + } + + if (typeof task.queue?.concurrencyLimit === "number") { + throw new ServiceValidationError( + `Task "${task.id}" declares both a queue concurrencyLimit and the concurrency option; use concurrency.` + ); + } + + for (const name of concurrency.limits ?? []) { + assertValidConcurrencyLimitName(name); + } + + if (concurrency.inline && task.queue?.name) { + if ((concurrency.limits ?? []).length > 1) { + throw new ServiceValidationError( + `Task "${task.id}": an inline limit on a shared queue uses a gate slot, so at most one named limit can be combined with it.` + ); + } + } + } +} + +const CONCURRENCY_LIMIT_NAME_MAX_LENGTH = 128 - CONCURRENCY_LIMIT_QUEUE_PREFIX.length; + +/** + * Named limits require a strict charset so distinct declared names can never merge + * onto one row after queue-name sanitization (which strips disallowed characters). + */ +function assertValidConcurrencyLimitName(name: string): void { + if (!new RegExp(`^[a-zA-Z0-9_-]{1,${CONCURRENCY_LIMIT_NAME_MAX_LENGTH}}$`).test(name)) { + throw new ServiceValidationError( + `Concurrency limit name "${name}" must be 1-${CONCURRENCY_LIMIT_NAME_MAX_LENGTH} characters using only letters, numbers, underscores and hyphens.` + ); + } +} + +function concurrencyLimitQueueName(limitName: string): string { + assertValidConcurrencyLimitName(limitName); + return `${CONCURRENCY_LIMIT_QUEUE_PREFIX}${limitName}`; +} + +/** User queue names may not claim the reserved limit/ namespace. */ +function assertNotReservedQueueName(name: string, context: string): void { + if (sanitizeQueueName(name).startsWith(CONCURRENCY_LIMIT_QUEUE_PREFIX)) { + throw new ServiceValidationError( + `${context}: queue names starting with "${CONCURRENCY_LIMIT_QUEUE_PREFIX}" are reserved for concurrency limits.` + ); + } +} + +/** + * Materializes the worker's declared named concurrency limits (plus any names tasks + * reference without declaring, created uncapped) as LIMIT-role TaskQueue rows. The + * columns mirror the declared shape exactly: perKey caps each key pool (and the + * keyless pool); total caps everything together via the group set. + */ +async function createWorkerConcurrencyLimits( + metadata: BackgroundWorkerMetadata, + worker: BackgroundWorker, + environment: AuthenticatedEnvironment, + prisma: PrismaClientOrTransaction +) { + const declared = new Map((metadata.concurrencyLimits ?? []).map((l) => [l.name, l])); + + for (const task of metadata.tasks) { + for (const name of task.concurrency?.limits ?? []) { + if (!declared.has(name)) { + declared.set(name, { name }); + } + } + } + + for (const limit of declared.values()) { + await createWorkerQueue( + { + name: concurrencyLimitQueueName(limit.name), + concurrencyLimit: limit.perKey ?? null, + combinedConcurrencyLimit: limit.total ?? null, + }, + limit.name, + "NAMED", + worker, + environment, + prisma, + "LIMIT", + "V2" + ); + } +} + async function createWorkerQueue( queue: QueueManifest, orderableName: string, queueType: TaskQueueType, worker: BackgroundWorker, environment: AuthenticatedEnvironment, - prisma: PrismaClientOrTransaction + prisma: PrismaClientOrTransaction, + role: TaskQueueRole = "QUEUE", + concurrencyVersion: TaskQueueConcurrencyVersion = "V1" ) { let queueName = sanitizeQueueName(queue.name); @@ -577,44 +903,110 @@ async function createWorkerQueue( const taskQueue = await upsertWorkerQueueRecord( queueName, baseConcurrencyLimit ?? null, + queue.combinedConcurrencyLimit ?? null, orderableName, queueType, worker, - prisma + prisma, + 0, + role, + concurrencyVersion ); - const newConcurrencyLimit = taskQueue.concurrencyLimit; + const syncQueueLimitsToEngine = async (row: { + name: string; + paused: boolean; + concurrencyLimit: number | null; + totalConcurrencyLimit: number | null; + }) => { + /** + * The total limit key is separate from the per-queue limit key that pause zeroes, + * so it is safe to sync it regardless of the paused state. The engine clamps it + * to the environment limit at read time, so the raw declared value is stored. + */ + if (typeof row.totalConcurrencyLimit === "number") { + await updateQueueTotalConcurrencyLimits(environment, row.name, row.totalConcurrencyLimit); + } else { + await removeQueueTotalConcurrencyLimits(environment, row.name); + } - if (!taskQueue.paused) { - if (typeof newConcurrencyLimit === "number") { - logger.debug("createWorkerQueue: updating concurrency limit", { + if (!row.paused) { + logger.debug("createWorkerQueue: syncing concurrency limit", { workerId: worker.id, - taskQueue, + taskQueue: row, orgId: environment.organizationId, projectId: environment.projectId, environmentId: environment.id, - concurrencyLimit: newConcurrencyLimit, + concurrencyLimit: row.concurrencyLimit, }); - await updateQueueConcurrencyLimits(environment, taskQueue.name, newConcurrencyLimit); + if (typeof row.concurrencyLimit === "number") { + await updateQueueConcurrencyLimits(environment, row.name, row.concurrencyLimit); + } else { + await removeQueueConcurrencyLimits(environment, row.name); + } } else { - logger.debug("createWorkerQueue: removing concurrency limit", { + /** + * A paused queue's engine limit is 0 (what pause wrote). Re-asserting it here + * heals the race where a pause lands between this deploy's row read and its + * engine sync, which would otherwise overwrite the 0 with the declared limit + * and leave a queue the dashboard shows as paused still dequeuing. + */ + logger.debug("createWorkerQueue: queue is paused, re-asserting the paused limit", { workerId: worker.id, - taskQueue, + taskQueue: row, orgId: environment.organizationId, projectId: environment.projectId, environmentId: environment.id, - concurrencyLimit: newConcurrencyLimit, }); - await removeQueueConcurrencyLimits(environment, taskQueue.name); + await updateQueueConcurrencyLimits(environment, row.name, 0); } - } else { - logger.debug("createWorkerQueue: queue is paused, not updating concurrency limit", { - workerId: worker.id, - taskQueue, - orgId: environment.organizationId, - projectId: environment.projectId, - environmentId: environment.id, + }; + + await syncQueueLimitsToEngine(taskQueue); + + /** + * The optimistic markers only guard the Postgres write; an override or reset can + * still land between that write and the engine sync above, which would leave the + * engine holding this deploy's stale values. Re-read the markers and re-sync from + * the fresh row until they stop moving (bounded): every actor writes Postgres + * before its own engine sync, so re-syncing whatever is freshest converges. A + * marker moving after the final read is healed by that actor's own engine sync + * or the next deploy. + */ + let syncedMarkers = { + concurrency: taskQueue.concurrencyLimitOverriddenAt?.getTime(), + total: taskQueue.totalConcurrencyLimitOverriddenAt?.getTime(), + paused: taskQueue.paused, + }; + + for (let i = 0; i < 3; i++) { + const freshQueue = await prisma.taskQueue.findFirst({ + where: { id: taskQueue.id }, + select: { + name: true, + paused: true, + concurrencyLimit: true, + totalConcurrencyLimit: true, + concurrencyLimitOverriddenAt: true, + totalConcurrencyLimitOverriddenAt: true, + }, }); + + if ( + !freshQueue || + (freshQueue.concurrencyLimitOverriddenAt?.getTime() === syncedMarkers.concurrency && + freshQueue.totalConcurrencyLimitOverriddenAt?.getTime() === syncedMarkers.total && + freshQueue.paused === syncedMarkers.paused) + ) { + break; + } + + await syncQueueLimitsToEngine(freshQueue); + syncedMarkers = { + concurrency: freshQueue.concurrencyLimitOverriddenAt?.getTime(), + total: freshQueue.totalConcurrencyLimitOverriddenAt?.getTime(), + paused: freshQueue.paused, + }; } return taskQueue; @@ -623,11 +1015,14 @@ async function createWorkerQueue( async function upsertWorkerQueueRecord( queueName: string, concurrencyLimit: number | null, + totalConcurrencyLimit: number | null, orderableName: string, queueType: TaskQueueType, worker: BackgroundWorker, prisma: PrismaClientOrTransaction, - attempt: number = 0 + attempt: number = 0, + role: TaskQueueRole = "QUEUE", + concurrencyVersion: TaskQueueConcurrencyVersion = "V1" ): Promise { if (attempt > 3) { throw new Error("Failed to insert queue record"); @@ -641,14 +1036,34 @@ async function upsertWorkerQueueRecord( }, }); + /** + * Queue names are unique per environment with no role distinction, and slashes were + * historically legal in user queue names, so a pre-existing user queue can occupy a + * limit/ name (or, defensively, the reverse). Silently reusing such a row would + * overwrite its limits and hide it from its management surface, so a role mismatch + * fails the deploy with an actionable rename instead. + */ + if (taskQueue && taskQueue.role !== role) { + throw new ServiceValidationError( + role === "LIMIT" + ? `Concurrency limit "${queueName.slice( + CONCURRENCY_LIMIT_QUEUE_PREFIX.length + )}" collides with an existing queue named "${queueName}" in this environment. Rename the concurrency limit and deploy again (renaming the queue in code does not free the name, because its historical queue row keeps it).` + : `Queue "${queueName}" collides with an existing concurrency limit row in this environment. Rename the queue and deploy again.` + ); + } + if (!taskQueue) { taskQueue = await prisma.taskQueue.create({ data: { friendlyId: generateFriendlyId("queue"), version: "V2", + role, + concurrencyVersion, name: queueName, orderableName, concurrencyLimit, + totalConcurrencyLimit, runtimeEnvironmentId: worker.runtimeEnvironmentId, projectId: worker.projectId, type: queueType, @@ -661,34 +1076,55 @@ async function upsertWorkerQueueRecord( }); } else { const hasOverride = taskQueue.concurrencyLimitOverriddenAt !== null; - + const hasTotalOverride = taskQueue.totalConcurrencyLimitOverriddenAt !== null; + + /** + * The override markers in the where clause make this an optimistic-concurrency + * update: a concurrent override/reset between the read above and this write + * changes a marker, the update misses (P2025) and the catch below retries with + * a fresh read, so a deploy can never clobber an operator's override. + */ taskQueue = await prisma.taskQueue.update({ where: { id: taskQueue.id, + concurrencyLimitOverriddenAt: taskQueue.concurrencyLimitOverriddenAt, + totalConcurrencyLimitOverriddenAt: taskQueue.totalConcurrencyLimitOverriddenAt, }, data: { workers: { connect: { id: worker.id } }, version: "V2", + concurrencyVersion, orderableName, // If overridden, keep current limit and update base; otherwise update limit normally concurrencyLimit: hasOverride ? undefined : concurrencyLimit, concurrencyLimitBase: hasOverride ? concurrencyLimit : undefined, + totalConcurrencyLimit: hasTotalOverride ? undefined : totalConcurrencyLimit, + totalConcurrencyLimitBase: hasTotalOverride ? totalConcurrencyLimit : undefined, }, }); } return taskQueue; } catch (error) { - // If the queue already exists, let's try again - if (error instanceof Prisma.PrismaClientKnownRequestError && error.code === "P2002") { + /** + * P2002: the queue was created concurrently. P2025: an override/reset moved a + * marker under the optimistic update. Both re-read and retry. + */ + if ( + error instanceof Prisma.PrismaClientKnownRequestError && + (error.code === "P2002" || error.code === "P2025") + ) { return await upsertWorkerQueueRecord( queueName, concurrencyLimit, + totalConcurrencyLimit, orderableName, queueType, worker, prisma, - attempt + 1 + attempt + 1, + role, + concurrencyVersion ); } throw error; diff --git a/apps/webapp/app/v3/services/createDeploymentBackgroundWorkerV4.server.ts b/apps/webapp/app/v3/services/createDeploymentBackgroundWorkerV4.server.ts index 5a003ac5198..185c1e41b89 100644 --- a/apps/webapp/app/v3/services/createDeploymentBackgroundWorkerV4.server.ts +++ b/apps/webapp/app/v3/services/createDeploymentBackgroundWorkerV4.server.ts @@ -15,6 +15,7 @@ import { syncDeclarativeSchedules, syncDeclarativeWebhooks, type BackgroundWorkerWithWarnings, + validateWorkerConcurrencyDeclarations, } from "./createBackgroundWorker.server"; import { findOrCreateBackgroundWorker } from "./createDeploymentBackgroundWorkerV4/findOrCreateBackgroundWorker.server"; import { TimeoutDeploymentService } from "./timeoutDeployment.server"; @@ -69,6 +70,23 @@ export class CreateDeploymentBackgroundWorkerServiceV4 extends BaseService { return; } + /** + * Reject invalid concurrency declarations before any worker rows exist. Queue + * rows and their engine limit keys are per-environment, so a mid-creation + * failure would leave the running version's limits already mutated. + */ + try { + validateWorkerConcurrencyDeclarations(body.metadata); + } catch (concurrencyError) { + if (concurrencyError instanceof ServiceValidationError) { + logger.warn("Invalid worker concurrency declarations", { + error: concurrencyError.message, + }); + await this.#failBackgroundWorkerDeployment(deployment, concurrencyError, environment); + } + throw concurrencyError; + } + // Handle multi-platform builds const deploymentPlatforms = deployment.imagePlatform?.split(",") ?? []; if (deploymentPlatforms.length > 1) { diff --git a/apps/webapp/app/v3/services/pauseQueue.server.ts b/apps/webapp/app/v3/services/pauseQueue.server.ts index aa3e21f9727..ad7b8f20e46 100644 --- a/apps/webapp/app/v3/services/pauseQueue.server.ts +++ b/apps/webapp/app/v3/services/pauseQueue.server.ts @@ -57,11 +57,20 @@ export class PauseQueueService extends BaseService { }, }); + /** + * Resume syncs from the row the update returned, never the pre-update + * snapshot, so a limit changed by a concurrent deploy is not resurrected. + * A declared limit of zero is a real limit and must be written, not removed. + */ if (action === "paused") { await updateQueueConcurrencyLimits(environment, queue.name, 0); } else { - if (queue.concurrencyLimit) { - await updateQueueConcurrencyLimits(environment, queue.name, queue.concurrencyLimit); + if (typeof updatedQueue.concurrencyLimit === "number") { + await updateQueueConcurrencyLimits( + environment, + queue.name, + updatedQueue.concurrencyLimit + ); } else { await removeQueueConcurrencyLimits(environment, queue.name); } @@ -85,6 +94,7 @@ export class PauseQueueService extends BaseService { friendlyId: updatedQueue.friendlyId, name: updatedQueue.name, type: updatedQueue.type, + version: updatedQueue.concurrencyVersion, running: results[1]?.[updatedQueue.name] ?? 0, queued: results[0]?.[updatedQueue.name] ?? 0, concurrencyLimit: updatedQueue.concurrencyLimit ?? null, diff --git a/apps/webapp/app/v3/services/replayTaskRun.server.ts b/apps/webapp/app/v3/services/replayTaskRun.server.ts index 4f1837e1002..c494f03b6b9 100644 --- a/apps/webapp/app/v3/services/replayTaskRun.server.ts +++ b/apps/webapp/app/v3/services/replayTaskRun.server.ts @@ -144,6 +144,15 @@ export class ReplayTaskRunService extends BaseService { : undefined, concurrencyKey: overrideOptions.concurrencyKey ?? existingTaskRun.concurrencyKey ?? undefined, + /** + * A run that held gates replays with those same gates. A run with none + * stored passes undefined (never a fabricated empty array, which reads + * as "clear the named limits") so the replay honors whatever limits the + * task declares now, like it honors the current queue and retry config. + */ + gates: Array.isArray(existingTaskRun.gates) + ? (existingTaskRun.gates as Array<{ queue: string; concurrencyKey?: string }>) + : undefined, maxAttempts: overrideOptions.maxAttempts, maxDuration: overrideOptions.maxDurationSeconds, machine: diff --git a/apps/webapp/app/v3/webhookEngine.server.ts b/apps/webapp/app/v3/webhookEngine.server.ts index b06f2f0024b..f45e15f8735 100644 --- a/apps/webapp/app/v3/webhookEngine.server.ts +++ b/apps/webapp/app/v3/webhookEngine.server.ts @@ -7,6 +7,7 @@ import { env } from "~/env.server"; import { findEnvironmentById } from "~/models/runtimeEnvironment.server"; import { logger } from "~/services/logger.server"; import { S2RealtimeStreams } from "~/services/realtime/s2realtimeStreams.server"; +import { SessionTriggerConfig as SessionTriggerConfigSchema } from "@trigger.dev/core/v3"; import { ensureRunForSession, type SessionTriggerConfig, @@ -154,17 +155,6 @@ function createWebhookEngine() { return { success: false, errorType: "NOT_FOUND", error: "Environment not found" }; } - const template = (triggerConfigTemplate ?? {}) as Partial; - const triggerConfig: SessionTriggerConfig = { - ...template, - basePayload: { - messages: [], - trigger: "preload", - chatId: externalId, - ...(template.basePayload ?? {}), - }, - }; - // Resume an existing session; otherwise only START one when the event is a session-start // (startOn). Resume-only with no session yet -> ignore (no session, no run, no egress). const existing = await findSessionByExternalId(environment, externalId); @@ -175,15 +165,64 @@ function createWebhookEngine() { skippedReason: "startOn: not a session-start event", }; } - const { session, isCached } = existing - ? { session: existing, isCached: true } - : await findOrCreateSession({ - environment, - externalId, - type: "chat.agent", - taskIdentifier, - triggerConfig, - }); + + let session; + let isCached; + if (existing) { + session = existing; + isCached = true; + } else { + /** The template arrives unvalidated (`z.record(z.unknown())` on the routing + * target), and continuations re-parse the stored row with a throwing parse — + * so anything this path persists must parse, or the session strands forever. + * A bad template fails the CREATE delivery terminally; resumes above never + * touch the template, so a broken template can't stop existing sessions. + * Known fields persist normalized (the parse output) while unknown template + * keys are kept as the pre-validation path stored them; a non-object + * `basePayload` is rejected rather than spread into index-keyed garbage. */ + const template = (triggerConfigTemplate ?? {}) as Partial; + if ( + template.basePayload !== undefined && + (typeof template.basePayload !== "object" || + template.basePayload === null || + Array.isArray(template.basePayload)) + ) { + return { + success: false, + error: + "Invalid triggerConfigTemplate on the webhook routing target: basePayload must be an object", + }; + } + const assembled = { + ...template, + basePayload: { + messages: [], + trigger: "preload", + chatId: externalId, + ...(template.basePayload ?? {}), + }, + }; + const parsedTriggerConfig = SessionTriggerConfigSchema.safeParse(assembled); + if (!parsedTriggerConfig.success) { + return { + success: false, + error: `Invalid triggerConfigTemplate on the webhook routing target: ${parsedTriggerConfig.error.issues + .map((issue) => `${issue.path.join(".")}: ${issue.message}`) + .join("; ")}`, + }; + } + const triggerConfig: SessionTriggerConfig = { + ...assembled, + ...parsedTriggerConfig.data, + }; + ({ session, isCached } = await findOrCreateSession({ + environment, + externalId, + type: "chat.agent", + taskIdentifier, + triggerConfig, + })); + } if (session.closedAt || (session.expiresAt && session.expiresAt.getTime() < Date.now())) { return { success: false, error: "Session is closed or expired" }; diff --git a/apps/webapp/test/concurrencyLimitsSystem.test.ts b/apps/webapp/test/concurrencyLimitsSystem.test.ts new file mode 100644 index 00000000000..a527678279b --- /dev/null +++ b/apps/webapp/test/concurrencyLimitsSystem.test.ts @@ -0,0 +1,564 @@ +import { postgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import { beforeEach, describe, expect, vi } from "vitest"; +import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; +import { anonymousConcurrencyLimitQueueName } from "~/v3/services/concurrencyLimitNames.server"; +import { ConcurrencyLimitsSystem } from "~/v3/services/concurrencyLimitsSystem.server"; + +/** + * These tests exercise the DB-write, marker and ordering logic against a real + * Postgres. The engine syncs are spies so tests can assert ordering and inject + * failures; the Redis side itself is covered by the run-engine suites. + */ +const { perKeySyncMock, perKeyRemoveMock, totalSyncMock, totalRemoveMock } = vi.hoisted(() => ({ + perKeySyncMock: vi.fn(async (..._args: unknown[]) => undefined), + perKeyRemoveMock: vi.fn(async (..._args: unknown[]) => undefined), + totalSyncMock: vi.fn(async (..._args: unknown[]) => undefined), + totalRemoveMock: vi.fn(async (..._args: unknown[]) => undefined), +})); + +vi.mock("~/v3/runQueue.server", () => ({ + updateQueueConcurrencyLimits: perKeySyncMock, + removeQueueConcurrencyLimits: perKeyRemoveMock, + updateQueueTotalConcurrencyLimits: totalSyncMock, + removeQueueTotalConcurrencyLimits: totalRemoveMock, +})); + +vi.mock("~/v3/runEngine.server", () => ({ + engine: { + totalConcurrencyOfQueues: async (_env: unknown, queues: string[]) => + Object.fromEntries(queues.map((q) => [q, 0])), + gateQueuedCountOfQueues: async (_env: unknown, queues: string[]) => + Object.fromEntries(queues.map((q) => [q, 0])), + currentConcurrencyOfQueues: async (_env: unknown, queues: string[]) => + Object.fromEntries(queues.map((q) => [q, 0])), + lengthOfQueues: async (_env: unknown, queues: string[]) => + Object.fromEntries(queues.map((q) => [q, 0])), + }, +})); + +vi.setConfig({ testTimeout: 30_000 }); + +async function seedEnvAndLimit( + prisma: PrismaClient, + opts: { perKey?: number | null; total?: number | null } = {} +) { + const slug = `s${Math.random().toString(36).slice(2, 10)}`; + + const organization = await prisma.organization.create({ data: { title: slug, slug } }); + const project = await prisma.project.create({ + data: { name: slug, slug, organizationId: organization.id, externalRef: slug }, + }); + const environment = await prisma.runtimeEnvironment.create({ + data: { + slug, + type: "PRODUCTION", + projectId: project.id, + organizationId: organization.id, + apiKey: slug, + pkApiKey: slug, + shortcode: slug, + maximumConcurrencyLimit: 100, + }, + }); + + const row = await prisma.taskQueue.create({ + data: { + friendlyId: `queue_${slug}`, + name: "limit/openai", + orderableName: "openai", + projectId: project.id, + runtimeEnvironmentId: environment.id, + role: "LIMIT", + concurrencyVersion: "V2", + concurrencyLimit: opts.perKey ?? null, + totalConcurrencyLimit: opts.total ?? null, + }, + }); + + const authEnv = { + id: environment.id, + maximumConcurrencyLimit: environment.maximumConcurrencyLimit, + } as unknown as AuthenticatedEnvironment; + + const system = new ConcurrencyLimitsSystem({ db: prisma, reader: prisma }); + + return { environment, row, authEnv, system }; +} + +describe("ConcurrencyLimitsSystem", () => { + /** Call counts must start at zero per test and leaked one-off implementations + * must not outlive the test that set them; mockReset also restores the default + * implementations given to vi.fn above. */ + beforeEach(() => { + perKeySyncMock.mockReset(); + perKeyRemoveMock.mockReset(); + totalSyncMock.mockReset(); + totalRemoveMock.mockReset(); + }); + + postgresTest( + "override changes only the given bound and keeps the declared base", + async ({ prisma }) => { + const { authEnv, system, row } = await seedEnvAndLimit(prisma, { total: 25 }); + + const result = await system.limits.override(authEnv, "openai", { total: 50 }); + expect(result.isOk()).toBe(true); + if (result.isOk()) { + expect(result.value.total).toMatchObject({ current: 50, base: 25, override: 50 }); + expect(result.value.perKey).toMatchObject({ current: null, override: null }); + } + + const updated = await prisma.taskQueue.findFirstOrThrow({ where: { id: row.id } }); + expect(updated.totalConcurrencyLimit).toBe(50); + expect(updated.totalConcurrencyLimitBase).toBe(25); + expect(updated.totalConcurrencyLimitOverriddenAt).not.toBeNull(); + expect(updated.concurrencyLimitOverriddenAt).toBeNull(); + + expect(totalSyncMock).toHaveBeenCalledWith(authEnv, "limit/openai", 50); + expect(perKeyRemoveMock).toHaveBeenCalledWith(authEnv, "limit/openai"); + } + ); + + postgresTest("override to zero pauses the limit in the engine", async ({ prisma }) => { + const { authEnv, system } = await seedEnvAndLimit(prisma, { total: 25 }); + + const result = await system.limits.override(authEnv, "openai", { total: 0 }); + expect(result.isOk()).toBe(true); + expect(totalSyncMock).toHaveBeenCalledWith(authEnv, "limit/openai", 0); + }); + + postgresTest("reset restores the declared values and clears the markers", async ({ prisma }) => { + const { authEnv, system, row } = await seedEnvAndLimit(prisma, { perKey: 2, total: 25 }); + + await system.limits.override(authEnv, "openai", { perKey: 10, total: 50 }); + const result = await system.limits.reset(authEnv, "openai"); + expect(result.isOk()).toBe(true); + if (result.isOk()) { + expect(result.value.total).toMatchObject({ current: 25, base: 25, override: null }); + expect(result.value.perKey).toMatchObject({ current: 2, base: 2, override: null }); + } + + const updated = await prisma.taskQueue.findFirstOrThrow({ where: { id: row.id } }); + expect(updated.concurrencyLimit).toBe(2); + expect(updated.totalConcurrencyLimit).toBe(25); + expect(updated.concurrencyLimitOverriddenAt).toBeNull(); + expect(updated.totalConcurrencyLimitOverriddenAt).toBeNull(); + }); + + postgresTest("reset without an override is rejected", async ({ prisma }) => { + const { authEnv, system } = await seedEnvAndLimit(prisma, { total: 25 }); + + const result = await system.limits.reset(authEnv, "openai"); + expect(result.isErr()).toBe(true); + if (result.isErr()) { + expect(result.error.type).toBe("limit_not_overridden"); + } + }); + + postgresTest( + "a failed engine sync during reset leaves the override intact so a retry converges", + async ({ prisma }) => { + const { authEnv, system, row } = await seedEnvAndLimit(prisma, { total: 25 }); + + await system.limits.override(authEnv, "openai", { total: 50 }); + + totalSyncMock.mockRejectedValueOnce(new Error("redis down")); + const failed = await system.limits.reset(authEnv, "openai"); + expect(failed.isErr()).toBe(true); + + /** The marker must survive the failed sync: the DB still says overridden. */ + const midway = await prisma.taskQueue.findFirstOrThrow({ where: { id: row.id } }); + expect(midway.totalConcurrencyLimitOverriddenAt).not.toBeNull(); + expect(midway.totalConcurrencyLimit).toBe(50); + + const retried = await system.limits.reset(authEnv, "openai"); + expect(retried.isOk()).toBe(true); + const final = await prisma.taskQueue.findFirstOrThrow({ where: { id: row.id } }); + expect(final.totalConcurrencyLimit).toBe(25); + expect(final.totalConcurrencyLimitOverriddenAt).toBeNull(); + } + ); + + postgresTest( + "a mutation whose markers moved underneath it conflicts instead of clobbering", + async ({ prisma }) => { + const { authEnv, system, row } = await seedEnvAndLimit(prisma, { total: 25 }); + + await system.limits.override(authEnv, "openai", { total: 50 }); + + /** + * Interleave a concurrent reset between this mutation's read and its write: + * the engine sync hook is the seam after the read, so clearing the markers + * there makes the guarded update miss and surface a conflict. + */ + totalSyncMock.mockImplementationOnce(async () => { + await prisma.taskQueue.update({ + where: { id: row.id }, + data: { + totalConcurrencyLimit: 25, + totalConcurrencyLimitBase: null, + totalConcurrencyLimitOverriddenAt: null, + totalConcurrencyLimitOverriddenBy: null, + }, + }); + }); + + const raced = await system.limits.reset(authEnv, "openai"); + expect(raced.isErr()).toBe(true); + if (raced.isErr()) { + expect(raced.error.type).toBe("conflict"); + } + + /** The concurrent actor's state stands untouched. */ + const final = await prisma.taskQueue.findFirstOrThrow({ where: { id: row.id } }); + expect(final.totalConcurrencyLimit).toBe(25); + expect(final.totalConcurrencyLimitOverriddenAt).toBeNull(); + } + ); + + postgresTest( + "a failed engine sync during override compensates from the fresh row", + async ({ prisma }) => { + const { authEnv, system, row } = await seedEnvAndLimit(prisma, { total: 25 }); + + totalSyncMock.mockRejectedValueOnce(new Error("redis down")); + const failed = await system.limits.override(authEnv, "openai", { total: 50 }); + expect(failed.isErr()).toBe(true); + if (failed.isErr()) { + expect(failed.error.type).toBe("sync_limit_to_engine_failed"); + } + + /** The persist already happened; compensation re-syncs it so the engine + * doesn't keep enforcing the old bound while the API reports the new one. + * Exactly two calls: the rejected primary sync, then the compensating + * re-sync from the fresh row — without compensation there is only one. */ + const updated = await prisma.taskQueue.findFirstOrThrow({ where: { id: row.id } }); + expect(updated.totalConcurrencyLimit).toBe(50); + expect(totalSyncMock).toHaveBeenCalledTimes(2); + expect(totalSyncMock).toHaveBeenLastCalledWith(authEnv, "limit/openai", 50); + } + ); + + postgresTest( + "an older override's engine write landing last is repaired by the freshness re-check", + async ({ prisma }) => { + const { authEnv, system, row } = await seedEnvAndLimit(prisma, { total: 25 }); + + /** `engineTotal` is written when a sync "lands", so landing order can differ + * from call order: the first override's write is delayed until a second, + * newer override has fully completed, then lands with the stale value. */ + let engineTotal: number | null = null; + let secondResult: Awaited> | undefined; + totalSyncMock.mockImplementation(async (_env, _name, value) => { + engineTotal = value as number; + }); + totalSyncMock.mockImplementationOnce(async (_env, _name, value) => { + secondResult = await system.limits.override(authEnv, "openai", { total: 75 }); + engineTotal = value as number; + }); + + const first = await system.limits.override(authEnv, "openai", { total: 50 }); + expect(first.isOk()).toBe(true); + expect(secondResult?.isOk()).toBe(true); + + const final = await prisma.taskQueue.findFirstOrThrow({ where: { id: row.id } }); + expect(final.totalConcurrencyLimit).toBe(75); + expect(engineTotal).toBe(75); + } + ); + + postgresTest( + "a default-queue inline limit resolves, overrides and resets under its task/ name", + async ({ prisma }) => { + const { authEnv, system, environment } = await seedEnvAndLimit(prisma, { total: 25 }); + + const queueRow = await prisma.taskQueue.create({ + data: { + friendlyId: `queue_t${environment.slug}`, + name: "task/send-email", + orderableName: "send-email", + projectId: environment.projectId, + runtimeEnvironmentId: environment.id, + role: "QUEUE", + concurrencyVersion: "V2", + concurrencyLimit: 1, + totalConcurrencyLimit: 10, + }, + }); + + const retrieved = await system.limits.retrieve(authEnv, "task/send-email"); + expect(retrieved.isOk()).toBe(true); + if (retrieved.isOk()) { + expect(retrieved.value.name).toBe("task/send-email"); + expect(retrieved.value.perKey).toMatchObject({ current: 1, base: 1 }); + expect(retrieved.value.total).toMatchObject({ current: 10, base: 10 }); + } + + const overridden = await system.limits.override(authEnv, "task/send-email", { total: 20 }); + expect(overridden.isOk()).toBe(true); + expect(totalSyncMock).toHaveBeenCalledWith(authEnv, "task/send-email", 20); + + const reset = await system.limits.reset(authEnv, "task/send-email"); + expect(reset.isOk()).toBe(true); + const final = await prisma.taskQueue.findFirstOrThrow({ where: { id: queueRow.id } }); + expect(final.totalConcurrencyLimit).toBe(10); + expect(final.totalConcurrencyLimitOverriddenAt).toBeNull(); + + const listed = await system.limits.list(authEnv, { page: 1, perPage: 50 }); + expect(listed.isOk()).toBe(true); + if (listed.isOk()) { + expect(listed.value.map((item) => item.name).sort()).toEqual(["openai", "task/send-email"]); + } + } + ); + + postgresTest( + "V1 queue rows and boundless V2 queue rows never surface as limits", + async ({ prisma }) => { + const { authEnv, system, environment } = await seedEnvAndLimit(prisma, { total: 25 }); + + await prisma.taskQueue.create({ + data: { + friendlyId: `queue_v1${environment.slug}`, + name: "task/legacy-task", + orderableName: "legacy-task", + projectId: environment.projectId, + runtimeEnvironmentId: environment.id, + role: "QUEUE", + concurrencyVersion: "V1", + concurrencyLimit: 5, + }, + }); + await prisma.taskQueue.create({ + data: { + friendlyId: `queue_nb${environment.slug}`, + name: "task/unbounded-task", + orderableName: "unbounded-task", + projectId: environment.projectId, + runtimeEnvironmentId: environment.id, + role: "QUEUE", + concurrencyVersion: "V2", + }, + }); + + const v1 = await system.limits.retrieve(authEnv, "task/legacy-task"); + expect(v1.isErr()).toBe(true); + + /** Boundless V2 queues stay out of the list but resolve by name, so an + * operator can still cap an undeclared task through this surface. */ + const boundless = await system.limits.retrieve(authEnv, "task/unbounded-task"); + expect(boundless.isOk()).toBe(true); + if (boundless.isOk()) { + expect(boundless.value.perKey.current).toBeNull(); + expect(boundless.value.total.current).toBeNull(); + } + + const listed = await system.limits.list(authEnv, { page: 1, perPage: 50 }); + expect(listed.isOk()).toBe(true); + if (listed.isOk()) { + expect(listed.value.map((item) => item.name)).toEqual(["openai"]); + } + } + ); + + postgresTest( + "a retired anonymous LIMIT row falls through to the live queue row", + async ({ prisma }) => { + const { authEnv, system, environment } = await seedEnvAndLimit(prisma, { total: 25 }); + + await prisma.taskQueue.create({ + data: { + friendlyId: `queue_rl${environment.slug}`, + name: "limit/task/send-email", + orderableName: "send-email", + projectId: environment.projectId, + runtimeEnvironmentId: environment.id, + role: "LIMIT", + concurrencyVersion: "V2", + }, + }); + await prisma.taskQueue.create({ + data: { + friendlyId: `queue_ql${environment.slug}`, + name: "task/send-email", + orderableName: "send-email-q", + projectId: environment.projectId, + runtimeEnvironmentId: environment.id, + role: "QUEUE", + concurrencyVersion: "V2", + concurrencyLimit: 1, + totalConcurrencyLimit: 10, + }, + }); + + const retrieved = await system.limits.retrieve(authEnv, "task/send-email"); + expect(retrieved.isOk()).toBe(true); + if (retrieved.isOk()) { + expect(retrieved.value.total).toMatchObject({ current: 10 }); + } + + const listed = await system.limits.list(authEnv, { page: 1, perPage: 50 }); + expect(listed.isOk()).toBe(true); + if (listed.isOk()) { + expect(listed.value.filter((item) => item.name === "task/send-email")).toHaveLength(1); + } + } + ); + + postgresTest("overrides and resets preserve a queue pause in the engine", async ({ prisma }) => { + const { authEnv, system, environment } = await seedEnvAndLimit(prisma, { total: 25 }); + + await prisma.taskQueue.create({ + data: { + friendlyId: `queue_p${environment.slug}`, + name: "task/paused-task", + orderableName: "paused-task", + projectId: environment.projectId, + runtimeEnvironmentId: environment.id, + role: "QUEUE", + concurrencyVersion: "V2", + concurrencyLimit: 1, + totalConcurrencyLimit: 10, + paused: true, + }, + }); + + const overridden = await system.limits.override(authEnv, "task/paused-task", { total: 20 }); + expect(overridden.isOk()).toBe(true); + expect(totalSyncMock).toHaveBeenCalledWith(authEnv, "task/paused-task", 20); + /** The pause IS the per-key engine value 0; the sync must rewrite 0, never + * the configured limit and never a removal. */ + expect(perKeySyncMock).toHaveBeenCalledWith(authEnv, "task/paused-task", 0); + expect(perKeySyncMock).not.toHaveBeenCalledWith(authEnv, "task/paused-task", 1); + expect(perKeyRemoveMock).not.toHaveBeenCalledWith(authEnv, "task/paused-task"); + + perKeySyncMock.mockClear(); + const reset = await system.limits.reset(authEnv, "task/paused-task"); + expect(reset.isOk()).toBe(true); + expect(perKeySyncMock).toHaveBeenCalledWith(authEnv, "task/paused-task", 0); + expect(perKeySyncMock).not.toHaveBeenCalledWith(authEnv, "task/paused-task", 1); + }); + + postgresTest("a perKey override clears a stale percent override source", async ({ prisma }) => { + const { authEnv, system, environment } = await seedEnvAndLimit(prisma, { total: 25 }); + + const row = await prisma.taskQueue.create({ + data: { + friendlyId: `queue_pc${environment.slug}`, + name: "task/percent-task", + orderableName: "percent-task", + projectId: environment.projectId, + runtimeEnvironmentId: environment.id, + role: "QUEUE", + concurrencyVersion: "V2", + concurrencyLimit: 50, + concurrencyLimitBase: 100, + concurrencyLimitOverriddenAt: new Date(), + concurrencyLimitOverridePercent: 50, + }, + }); + + const overridden = await system.limits.override(authEnv, "task/percent-task", { perKey: 3 }); + expect(overridden.isOk()).toBe(true); + const updated = await prisma.taskQueue.findFirstOrThrow({ where: { id: row.id } }); + expect(updated.concurrencyLimit).toBe(3); + expect(updated.concurrencyLimitOverridePercent).toBeNull(); + }); + + postgresTest("an uncapped named limit stays visible and cappable", async ({ prisma }) => { + const { authEnv, system, environment } = await seedEnvAndLimit(prisma, { total: 25 }); + + /** Deploys materialize referenced-but-undeclared names as uncapped LIMIT + * rows; only the anonymous limit/task/ namespace treats boundless as + * retired. */ + await prisma.taskQueue.create({ + data: { + friendlyId: `queue_un${environment.slug}`, + name: "limit/acme-api", + orderableName: "acme-api", + projectId: environment.projectId, + runtimeEnvironmentId: environment.id, + role: "LIMIT", + concurrencyVersion: "V2", + }, + }); + + const retrieved = await system.limits.retrieve(authEnv, "acme-api"); + expect(retrieved.isOk()).toBe(true); + if (retrieved.isOk()) { + expect(retrieved.value.total.current).toBeNull(); + } + + const capped = await system.limits.override(authEnv, "acme-api", { total: 5 }); + expect(capped.isOk()).toBe(true); + expect(totalSyncMock).toHaveBeenCalledWith(authEnv, "limit/acme-api", 5); + + const listed = await system.limits.list(authEnv, { page: 1, perPage: 50 }); + expect(listed.isOk()).toBe(true); + if (listed.isOk()) { + expect(listed.value.map((item) => item.name)).toContain("acme-api"); + } + }); + + postgresTest("retrieve misses queue-role rows and unknown names", async ({ prisma }) => { + const { authEnv, system, environment } = await seedEnvAndLimit(prisma, { total: 25 }); + + await prisma.taskQueue.create({ + data: { + friendlyId: `queue_q${environment.slug}`, + name: "limit/shadow", + orderableName: "shadow", + projectId: environment.projectId, + runtimeEnvironmentId: environment.id, + role: "QUEUE", + }, + }); + + const missing = await system.limits.retrieve(authEnv, "missing"); + expect(missing.isErr()).toBe(true); + + const shadow = await system.limits.retrieve(authEnv, "shadow"); + expect(shadow.isErr()).toBe(true); + }); + + postgresTest( + "an anonymous task limit with a hashed row name resolves by its documented task/ name", + async ({ prisma }) => { + const { environment, system, authEnv } = await seedEnvAndLimit(prisma, {}); + const rowName = anonymousConcurrencyLimitQueueName("billing.charge"); + expect(rowName).not.toBe("limit/task/billing.charge"); + await prisma.taskQueue.create({ + data: { + friendlyId: `queue_${environment.id}_anon`, + name: rowName, + orderableName: rowName, + projectId: environment.projectId, + runtimeEnvironmentId: environment.id, + role: "LIMIT", + concurrencyVersion: "V2", + totalConcurrencyLimit: 3, + }, + }); + + const result = await system.limits.retrieve(authEnv, "task/billing.charge"); + expect(result.isOk()).toBe(true); + if (result.isOk()) { + expect(result.value.total).toMatchObject({ current: 3 }); + } + } + ); + + postgresTest( + "current clamps to the environment limit while base keeps the declared value", + async ({ prisma }) => { + const { system, authEnv } = await seedEnvAndLimit(prisma, { total: 200 }); + + const result = await system.limits.retrieve(authEnv, "openai"); + expect(result.isOk()).toBe(true); + if (result.isOk()) { + expect(result.value.total).toMatchObject({ current: 100, base: 200, override: null }); + } + } + ); +}); diff --git a/apps/webapp/test/concurrencySystemTotalResetConflict.test.ts b/apps/webapp/test/concurrencySystemTotalResetConflict.test.ts new file mode 100644 index 00000000000..64fb47a0ffb --- /dev/null +++ b/apps/webapp/test/concurrencySystemTotalResetConflict.test.ts @@ -0,0 +1,138 @@ +import { postgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import { beforeEach, describe, expect, vi } from "vitest"; +import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; +import { ConcurrencySystem } from "~/v3/services/concurrencySystem.server"; + +/** + * The reset flow syncs the engine to the declared base BEFORE its guarded DB write, so + * a concurrent override landing between the reset's read and its write both trips the + * optimistic guard AND leaves the reset's stale engine write behind. The mocked sync is + * the injection point: performing the concurrent override inside the enforce-first sync + * makes the race deterministic, and the assertions pin the heal that must follow. + */ +const { totalSyncMock, totalRemoveMock } = vi.hoisted(() => ({ + totalSyncMock: vi.fn(async (..._args: unknown[]) => undefined), + totalRemoveMock: vi.fn(async (..._args: unknown[]) => undefined), +})); + +vi.mock("~/v3/runEngine.server", () => ({ + engine: { + lengthOfQueues: async (_env: unknown, queues: string[]) => + Object.fromEntries(queues.map((q) => [q, 0])), + currentConcurrencyOfQueues: async (_env: unknown, queues: string[]) => + Object.fromEntries(queues.map((q) => [q, 0])), + runQueue: { + updateQueueConcurrencyLimits: async () => undefined, + removeQueueConcurrencyLimits: async () => undefined, + updateQueueTotalConcurrencyLimits: totalSyncMock, + removeQueueTotalConcurrencyLimits: totalRemoveMock, + updateEnvConcurrencyLimits: async () => undefined, + }, + }, +})); + +vi.setConfig({ testTimeout: 30_000 }); + +async function seedOverriddenQueue(prisma: PrismaClient) { + const slug = `s${Math.random().toString(36).slice(2, 10)}`; + const organization = await prisma.organization.create({ data: { title: slug, slug } }); + const project = await prisma.project.create({ + data: { name: slug, slug, organizationId: organization.id, externalRef: slug }, + }); + const environment = await prisma.runtimeEnvironment.create({ + data: { + slug, + type: "PRODUCTION", + projectId: project.id, + organizationId: organization.id, + apiKey: slug, + pkApiKey: slug, + shortcode: slug, + maximumConcurrencyLimit: 100, + }, + }); + const queue = await prisma.taskQueue.create({ + data: { + friendlyId: `queue_${slug}`, + name: "my-queue", + type: "NAMED", + version: "V2", + projectId: project.id, + runtimeEnvironmentId: environment.id, + totalConcurrencyLimit: 10, + totalConcurrencyLimitBase: 5, + totalConcurrencyLimitOverriddenAt: new Date(), + }, + }); + + const authEnv = { + id: environment.id, + maximumConcurrencyLimit: environment.maximumConcurrencyLimit, + } as unknown as AuthenticatedEnvironment; + const system = new ConcurrencySystem({ db: prisma, reader: prisma }); + + return { queue, authEnv, system }; +} + +describe("total-limit reset losing to a concurrent override", () => { + beforeEach(() => { + totalSyncMock.mockReset(); + totalRemoveMock.mockReset(); + }); + + postgresTest( + "returns a conflict and re-syncs the engine from the winning row", + async ({ prisma }) => { + const { queue, authEnv, system } = await seedOverriddenQueue(prisma); + + /** The enforce-first sync (base = 5) is where the concurrent override lands: it moves + * the marker so the reset's guarded write misses, and its own engine sync never runs + * (the winner's sync failing is the case that makes healing load-bearing). */ + totalSyncMock + .mockImplementationOnce(async () => { + await prisma.taskQueue.update({ + where: { id: queue.id }, + data: { + totalConcurrencyLimit: 20, + totalConcurrencyLimitBase: 5, + totalConcurrencyLimitOverriddenAt: new Date(), + }, + }); + }) + /** A THIRD writer landing while the heal's own engine write is in flight: the + * heal's convergence loop must re-read and re-apply it, so its stale write can + * never be the last one standing. */ + .mockImplementationOnce(async () => { + await prisma.taskQueue.update({ + where: { id: queue.id }, + data: { + totalConcurrencyLimit: 7, + totalConcurrencyLimitBase: 5, + totalConcurrencyLimitOverriddenAt: new Date(), + }, + }); + }); + + const result = await system.queues.resetTotalConcurrencyLimit(authEnv, queue.friendlyId); + + expect(result.isErr()).toBe(true); + if (result.isErr()) { + expect(result.error.type).toBe("concurrent_modification"); + } + + const row = await prisma.taskQueue.findFirstOrThrow({ where: { id: queue.id } }); + expect(row.totalConcurrencyLimit).toBe(7); + expect(row.totalConcurrencyLimitOverriddenAt).not.toBeNull(); + + /** Call 1 is the reset's enforce-first write of the stale base (5); the heal follows + * with 20 (already stale by the time it lands) and must then converge on 7, the + * persisted winner, as the LAST engine write. */ + const syncedValues = totalSyncMock.mock.calls.map((call) => call[2]); + expect(syncedValues[0]).toBe(5); + expect(syncedValues).toContain(20); + expect(syncedValues[syncedValues.length - 1]).toBe(7); + expect(totalRemoveMock).not.toHaveBeenCalled(); + } + ); +}); diff --git a/apps/webapp/test/dashboardAgentToolScopes.test.ts b/apps/webapp/test/dashboardAgentToolScopes.test.ts index 11236aae43b..c8276896f53 100644 --- a/apps/webapp/test/dashboardAgentToolScopes.test.ts +++ b/apps/webapp/test/dashboardAgentToolScopes.test.ts @@ -34,7 +34,7 @@ const VIA_ENV_JWT: Read[] = [ { tool: "get_queue (metrics)", path: "/api/v1/queues/:name/metrics", - resource: { type: "query", id: "queue_metrics" }, + resource: { type: "query", id: "concurrency_metrics" }, }, { tool: "get_queue (live row)", path: "/api/v1/queues/:name", resource: { type: "queues" } }, { diff --git a/apps/webapp/test/engine/triggerTask.metadataCache.test.ts b/apps/webapp/test/engine/triggerTask.metadataCache.test.ts index 1309f36d9f6..07ecf31ba41 100644 --- a/apps/webapp/test/engine/triggerTask.metadataCache.test.ts +++ b/apps/webapp/test/engine/triggerTask.metadataCache.test.ts @@ -73,6 +73,7 @@ describe("DefaultQueueManager task metadata cache", () => { slug: taskIdentifier, ttl: null, triggerSource: "AGENT", + gates: null, queueId: null, queueName: `task/${taskIdentifier}`, }, @@ -209,6 +210,7 @@ describe("DefaultQueueManager task metadata cache", () => { slug: taskIdentifier, ttl: null, triggerSource: "AGENT", + gates: null, queueId: null, queueName: `task/${taskIdentifier}`, }, @@ -286,6 +288,7 @@ describe("DefaultQueueManager task metadata cache", () => { slug: taskIdentifier, ttl: null, triggerSource: "AGENT", + gates: null, queueId: null, queueName: `task/${taskIdentifier}`, }, @@ -295,6 +298,7 @@ describe("DefaultQueueManager task metadata cache", () => { slug: taskIdentifier, ttl: null, triggerSource: "SCHEDULED", + gates: null, queueId: null, queueName: `task/${taskIdentifier}`, }, diff --git a/apps/webapp/test/reportHealth.test.ts b/apps/webapp/test/reportHealth.test.ts index 706291816f0..6357ac2f77b 100644 --- a/apps/webapp/test/reportHealth.test.ts +++ b/apps/webapp/test/reportHealth.test.ts @@ -94,8 +94,8 @@ describe("health cause tree (Golden A — env limit saturation)", () => { it("footer = raise the limit (self-serve) + docs + do-nothing (drains)", () => { expect(vm.footer).toEqual([ - { code: "raise_env_limit", link: "concurrency" }, - { code: "concurrency_docs", link: "concurrency" }, + { code: "raise_env_limit", link: "concurrency-limits" }, + { code: "concurrency_docs", link: "concurrency-limits" }, { code: "do_nothing_drains", value: 2.3 }, ]); }); diff --git a/apps/webapp/test/reportHealthData.test.ts b/apps/webapp/test/reportHealthData.test.ts index 6b12f33ea75..4beaf6546aa 100644 --- a/apps/webapp/test/reportHealthData.test.ts +++ b/apps/webapp/test/reportHealthData.test.ts @@ -51,7 +51,7 @@ function makeDeps(opts: { if (query.includes("dlq_total")) return wrap(opts.queueTotals); if (isEnv && query.includes("timeBucket")) return wrap(opts.envSeries); if (isEnv) return wrap(opts.envScalar ?? [{}]); - if (query.includes("FROM queue_metrics")) return wrap(opts.worst); + if (query.includes("FROM concurrency_metrics")) return wrap(opts.worst); if (query.includes("task_identifier")) return wrap([]); if (query.includes("FROM runs") && query.includes("timeBucket")) return wrap(opts.runsSeries); return wrap(opts.runs); @@ -76,7 +76,7 @@ const RUNS_SCALAR: Rows = [ ]; describe("loadHealthInput — orchestration (query seam)", () => { - it("measured path: queue_metrics source, real pending, parsed dlq, window from timeRange", async () => { + it("measured path: concurrency_metrics source, real pending, parsed dlq, window from timeRange", async () => { const input = await loadHealthInput( fakeEnv, "1h", diff --git a/apps/webapp/test/reportsApiRoute.test.ts b/apps/webapp/test/reportsApiRoute.test.ts index f027c0d9eb3..37043b6513b 100644 --- a/apps/webapp/test/reportsApiRoute.test.ts +++ b/apps/webapp/test/reportsApiRoute.test.ts @@ -105,7 +105,7 @@ describe("api.v1.reports.$key — authorization", () => { expect(requiredResources("health")).toEqual([ { type: "query", id: "runs" }, { type: "query", id: "env_metrics" }, - { type: "query", id: "queue_metrics" }, + { type: "query", id: "concurrency_metrics" }, ]); }); @@ -118,7 +118,7 @@ describe("api.v1.reports.$key — authorization", () => { describe("reportQueryTables — scope derivation from the registry", () => { const registry: Record = { - health: { tables: ["runs", "env_metrics", "queue_metrics"] }, + health: { tables: ["runs", "env_metrics", "concurrency_metrics"] }, narrow: { tables: ["runs"] }, }; @@ -127,7 +127,11 @@ describe("reportQueryTables — scope derivation from the registry", () => { }); it("still gives the wider report all of its tables", () => { - expect(reportQueryTables("health", registry)).toEqual(["runs", "env_metrics", "queue_metrics"]); + expect(reportQueryTables("health", registry)).toEqual([ + "runs", + "env_metrics", + "concurrency_metrics", + ]); }); it("returns no tables for an unknown key", () => { diff --git a/apps/webapp/test/resolveTriggerUri.test.ts b/apps/webapp/test/resolveTriggerUri.test.ts index 3aaa9ae04b0..383bdba00d9 100644 --- a/apps/webapp/test/resolveTriggerUri.test.ts +++ b/apps/webapp/test/resolveTriggerUri.test.ts @@ -19,7 +19,7 @@ import { resolveTriggerUrisInOrganization } from "~/services/resolveTriggerUriIn import { v3DeploymentVersionPath, v3ErrorPath, - v3QueuesPath, + concurrencyPath, v3RunPath, v3RunSpanPath, } from "~/utils/pathBuilder"; @@ -73,7 +73,7 @@ describe("resolveTriggerUri", () => { const uri = formatTriggerUri({ kind: "queue", ...uriScope, name: "task/send email" }); expect(resolveTriggerUri(scope, uri)).toEqual({ label: "task/send email", - url: `${v3QueuesPath(org, project, env)}?query=task%2Fsend%20email`, + url: `${concurrencyPath(org, project, env)}?query=task%2Fsend%20email`, }); }); diff --git a/docs/ai-chat/client-protocol.mdx b/docs/ai-chat/client-protocol.mdx index b5007b0e52c..ec6becf3cd4 100644 --- a/docs/ai-chat/client-protocol.mdx +++ b/docs/ai-chat/client-protocol.mdx @@ -206,6 +206,8 @@ Pick `"preload"` when the UI has rendered but the user hasn't typed (warms the a | `expiresAt` | `string` (ISO date) | Retention cap. | | `triggerConfig.machine` | `string` | Machine preset (`micro`, `small-1x`, …) for every run. | | `triggerConfig.queue` | `string` | Queue name. | +| `triggerConfig.concurrency` | `string[]` | Up to two [named concurrency limits](/concurrency#sharing-a-limit-between-tasks) every run holds, replacing the task's declared named limits. | +| `triggerConfig.concurrencyKey` | `string` | Scopes every run of this session to its own pool under each `perKey` bound it holds. Never defaulted — pass one (e.g. your chat or tenant ID) to isolate sessions from each other. | | `triggerConfig.tags` | `string[]` | Tags applied to every run (in addition to session-level `tags`). | | `triggerConfig.maxAttempts` | `number` | Per-run retry cap (1–10). | | `triggerConfig.maxDuration` | `number` | Per-run wall-clock cap, seconds. | diff --git a/docs/concurrency.mdx b/docs/concurrency.mdx new file mode 100644 index 00000000000..02f17a8c055 --- /dev/null +++ b/docs/concurrency.mdx @@ -0,0 +1,370 @@ +--- +title: "Concurrency" +description: "Limit how many runs execute at once: per task, per tenant, or shared across tasks." +--- + +When you trigger a task, the [run](/runs) waits in a [queue](/queues) and runs start in trigger order as capacity allows. Concurrency is what decides how much capacity there is. + +By default, concurrency is only limited by your environment concurrency limit. If you need more control (for example, to limit concurrency or share limits across multiple tasks), you can set a concurrency limit on the task, or declare a named limit and share it, as described below. + +Controlling concurrency is useful when you have a task that can't be run concurrently, or when you want to limit the number of runs to avoid overloading a resource. + +It's important to note that only actively executing runs count towards concurrency limits. Runs that are delayed or waiting in a queue do not consume concurrency slots until they begin execution. + +## Use cases + +- **Limit how many runs of a task execute at once**: [Setting task concurrency](#setting-task-concurrency) +- **Share one limit across several tasks**: [Sharing a limit between tasks](#sharing-a-limit-between-tasks) +- **Give each tenant its own separate concurrency**: [Concurrency keys and per-tenant limits](#concurrency-keys-and-per-tenant-limits) +- **Per-tenant limits with a ceiling on the total**: [Per-key and total limits together](#per-key-and-total-limits-together) +- **Cap a tenant across every task they run**: [A per-tenant cap across multiple tasks](#a-per-tenant-cap-across-multiple-tasks) +- **Cap a shared resource, like an external API, across tasks and tenants**: [A global cap for a shared resource](#a-global-cap-for-a-shared-resource) +- **Switch a run's limits when you trigger it**: [Setting limits when you trigger a run](#setting-limits-when-you-trigger-a-run) + +## Default concurrency + +By default, all tasks have an unbounded concurrency limit, limited only by the overall concurrency limits of your environment. + + + Your environment has a base concurrency limit and a burstable limit (default burst factor of 2.0x + the base limit). Individual tasks and limits are capped by the base concurrency limit, not the + burstable limit. For example, if your base limit is 10, your environment can burst up to 20 + concurrent runs, but any single limit can allow at most 10 concurrent runs. If you're a paying + customer you can request higher burst limits by [contacting us](https://www.trigger.dev/contact). + + +## Setting task concurrency + +Set the `concurrency` option on a task to limit how many of its runs execute at once. `{ total: n }` caps the task outright: + +```ts /trigger/one-at-a-time.ts +import { task } from "@trigger.dev/sdk"; + +// This task will only run one at a time +export const oneAtATime = task({ + id: "one-at-a-time", + concurrency: { total: 1 }, + run: async (payload) => { + //... + }, +}); +``` + +This is useful if you need to control access to a shared resource, like a database or an API that has rate limits. + +There are two ways to bound a task, and you can combine them: + +- `total` caps every run of the task together, whether or not runs use a `concurrencyKey`. +- `perKey` caps each `concurrencyKey` pool separately; runs triggered without a key share one pool. + +```ts /trigger/per-user.ts +import { task } from "@trigger.dev/sdk"; + +export const generateReport = task({ + id: "generate-report", + // each user runs at most 1 at a time, and at most 10 run in total + concurrency: { perKey: 1, total: 10 }, + run: async (payload) => { + //... + }, +}); +``` + + + The `queue: {"{ concurrencyLimit: n }"}` option keeps working but is deprecated in favor of + `concurrency`. Its single number means "per key when runs pass a `concurrencyKey`, whole queue + when they don't" — the `concurrency` shape says which you mean explicitly. + + +## Sharing a limit between tasks + +Declare a named limit with `concurrencyLimit()` and put it in each task's `concurrency`. Every task holding the limit draws from the same pools: + +```ts /trigger/limits.ts +import { concurrencyLimit, task } from "@trigger.dev/sdk"; + +// at most 25 concurrent runs across every task that holds this limit +export const openaiLimit = concurrencyLimit({ name: "openai", total: 25 }); + +export const generateSummary = task({ + id: "generate-summary", + concurrency: openaiLimit, + run: async (payload) => { + // ... + }, +}); + +export const generateTitle = task({ + id: "generate-title", + concurrency: openaiLimit, + run: async (payload) => { + // ... + }, +}); +``` + +A task's `concurrency` takes a single item or an array: at most one inline shape (which caps that task alone) plus up to two named limits. A run starts only when every limit it holds has capacity, and it occupies a slot in each while it executes: + +```ts /trigger/summarize.ts +import { concurrencyLimit, task } from "@trigger.dev/sdk"; + +export const openaiLimit = concurrencyLimit({ name: "openai", total: 25 }); + +export const summarizeThread = task({ + id: "summarize-thread", + // this task runs at most 5 at once, and also counts towards the shared openai limit + concurrency: [{ total: 5 }, openaiLimit], + run: async (payload) => { + // ... + }, +}); +``` + + + Names you declare with `concurrencyLimit()` are 1-122 characters using only letters, numbers, + underscores and hyphens. + + +## Setting limits when you trigger a run + +The trigger-time `concurrency` option takes limit names and replaces the task's declared **named** limits for that run. The task's inline limit always applies: + +```ts app/api/report/route.ts +import { generateReport } from "~/trigger/reports"; + +// this run counts towards "priority" instead of the task's declared named limits +await generateReport.trigger(data, { concurrency: ["priority"] }); +``` + +Pass an empty array to run with only the task's inline limit. + +## Concurrency keys and per-tenant limits + +If you're building an application where you want to run tasks for your users, you might want a separate limit for each of your users (or orgs, projects, etc.). + +You can do this by passing a `concurrencyKey` when you trigger. Each unique key value gets its own pool under every `perKey` bound the run holds: + +```ts app/api/report/route.ts +import { generateReport } from "~/trigger/per-user"; + +export async function POST(request: Request) { + const data = await request.json(); + + const handle = await generateReport.trigger(data, { + // every user gets their own concurrency pool + concurrencyKey: data.userId, + }); + + return Response.json(handle); +} +``` + +## Per-key and total limits together + +`perKey` on its own lets total concurrency grow with the number of active keys: ten active users under `perKey: 5` can run 50 at once. Add `total` to bound everything as a group. Each key still gets at most `perKey`, and all runs together — keyed or not — never exceed `total`: + +```ts /trigger/per-user-capped.ts +import { task } from "@trigger.dev/sdk"; + +export const processUpload = task({ + id: "process-upload", + // each user runs at most 1 at a time, and at most 10 run in total + concurrency: { perKey: 1, total: 10 }, + run: async (payload) => { + //... + }, +}); +``` + +## A per-tenant cap across multiple tasks + +A named limit's `perKey` bound follows each run's own `concurrencyKey`, so one declaration caps each tenant across every task holding the limit: + +```ts /trigger/webhooks.ts +import { concurrencyLimit, task } from "@trigger.dev/sdk"; + +// each tenant runs at most 10 at once across every task that holds this limit +export const tenantLimit = concurrencyLimit({ name: "tenant", perKey: 10 }); + +export const processWebhook = task({ + id: "process-webhook", + // webhooks themselves are capped at 2 per tenant, within the tenant's overall 10 + concurrency: [{ perKey: 2 }, tenantLimit], + run: async (payload) => { + //... + }, +}); +``` + +```ts app/api/webhook/route.ts +// the run counts towards this tenant's pool in both limits +await processWebhook.trigger(payload, { concurrencyKey: tenantId }); +``` + +## A global cap for a shared resource + +To cap something global, like total traffic to an external API, across many tasks and all tenants: declare a limit with only a `total` and share it. It counts every run holding it, whether or not the run has a `concurrencyKey`: + +```ts /trigger/sync.ts +import { concurrencyLimit, task } from "@trigger.dev/sdk"; + +// at most 10 concurrent provider calls across every task and every tenant +export const providerApiLimit = concurrencyLimit({ name: "provider-api", total: 10 }); + +export const syncToProvider = task({ + id: "sync-to-provider", + concurrency: [{ total: 20 }, providerApiLimit], + run: async (payload) => { + //... + }, +}); +``` + +Runs with and without a `concurrencyKey` share the same `total`, so this works even when only some of your triggers have a natural key. + +## Concurrency and subtasks + +When you trigger a task that has subtasks, the subtasks will not inherit the parent's limits. Unless otherwise specified, subtasks run under their own task's configuration: + +```ts /trigger/subtasks.ts +export const parentTask = task({ + id: "parent-task", + run: async (payload) => { + //trigger a subtask + await subtask.triggerAndWait(payload); + }, +}); + +// This subtask runs under its own limits +export const subtask = task({ + id: "subtask", + run: async (payload) => { + //... + }, +}); +``` + +## Waits and concurrency + +With our [task checkpoint system](/how-it-works#the-checkpoint-resume-system), tasks can wait at various waitpoints (like waiting for subtasks to complete, delays, or external events). The way this system interacts with the concurrency system is important to understand. + +Concurrency is only released when a run reaches a waitpoint and is checkpointed. When a run is checkpointed, it transitions to the `WAITING` state and releases its concurrency slots back to every limit it holds and the environment, allowing other runs to execute or resume. + +This means that: + +- Only actively executing runs count towards concurrency limits +- Runs in the `WAITING` state (checkpointed at waitpoints) do not consume concurrency slots +- You can have more runs in the `WAITING` state than a limit allows to execute +- When a waiting run resumes (e.g., when a subtask completes), it must re-acquire its slots + +For example, if a task has `concurrency: { total: 1 }`: + +- You can only have exactly 1 run executing at a time +- You may have multiple runs in the `WAITING` state for that task +- When the executing run reaches a waitpoint and checkpoints, it releases its slot +- The next queued run can then begin execution + +### Short time-based waits keep their slot + +Checkpointing takes time, so a run doesn't checkpoint the moment it reaches a waitpoint. For [`wait.for()`](/wait-for) and [`wait.until()`](/wait-until) it happens 60 seconds into the wait, so anything shorter stays `EXECUTING` and holds its slots for the whole wait. If you're polling in a loop, use an interval comfortably above 60 seconds so the slots are actually released between polls. + +### Waiting for a subtask + +When a parent task triggers and waits for a subtask, the parent task will checkpoint and release its concurrency slots once it reaches the wait point. This prevents environment deadlocks where all concurrency slots would be occupied by waiting tasks. + +```ts /trigger/waiting.ts +export const parentTask = task({ + id: "parent-task", + concurrency: { total: 1 }, + run: async (payload) => { + //trigger a subtask and wait for it to complete + await subtask.triggerAndWait(payload); + // The parent task checkpoints here and releases its concurrency slots + // allowing other tasks to execute while waiting + }, +}); + +export const subtask = task({ + id: "subtask", + run: async (payload) => { + //... + }, +}); +``` + +When the parent task reaches the `triggerAndWait` call, it checkpoints and transitions to the `WAITING` state, releasing its slots. Once the subtask completes, the parent task will resume and re-acquire them. + +## Managing concurrency limits with the SDK + +The `concurrencyLimits` namespace manages your named limits at runtime (anonymous inline limits appear under derived `task/` names): + +```ts +import { concurrencyLimits } from "@trigger.dev/sdk"; +``` + +### Listing limits + +```ts +import { concurrencyLimits } from "@trigger.dev/sdk"; + +// List all limits (returns paginated results) +const allLimits = await concurrencyLimits.list(); + +// With pagination options +const pagedLimits = await concurrencyLimits.list({ + page: 1, + perPage: 20, +}); +``` + +### Retrieving a limit + +Retrieve a limit by its name to see its bounds and live counts: + +```ts +import { concurrencyLimits } from "@trigger.dev/sdk"; + +const limit = await concurrencyLimits.retrieve("openai"); +``` + +The limit object contains each bound plus live counts: + +```ts +{ + id: "climit_1234", + name: "openai", + perKey: { + current: null, // Enforced right now (null = no per-key bound) + base: null, // Declared in your code + override: null, // Override value (if set) + overriddenAt: null, // When the override was applied + }, + total: { + current: 25, + base: 25, + override: null, + overriddenAt: null, + }, + running: 14, // Runs executing that hold this limit + queued: 100, // Runs queued that must clear this limit to execute +} +``` + +### Overriding a limit + +Overrides change only the fields you pass; the declared values are kept and restored by `reset`. Overriding `total` to `0` blocks every run holding the limit, which is how you pause a limit: + +```ts +import { concurrencyLimits } from "@trigger.dev/sdk"; + +// Raise the shared cap +await concurrencyLimits.override("openai", { total: 50 }); + +// Pause the limit: nothing holding it can start +await concurrencyLimits.override("openai", { total: 0 }); + +// Back to the values declared in your code +await concurrencyLimits.reset("openai"); +``` + +Overrides survive deploys: redeploying your code keeps an active override until you reset it. diff --git a/docs/database-connections.mdx b/docs/database-connections.mdx index df808e66088..caf6b0927fa 100644 --- a/docs/database-connections.mdx +++ b/docs/database-connections.mdx @@ -68,7 +68,7 @@ Set the pool small. A task usually runs its queries in sequence, so one connecti | Drizzle (node-postgres) | 10 (the underlying `pg` pool) | | [MongoDB driver](https://www.mongodb.com/docs/drivers/node/current/connect/connection-options/connection-pools/) | 100 (`maxPoolSize`) | -Keep `concurrent runs × pool size` under your provider's connection limit, and cap how many runs execute at once with [concurrency limits](/queue-concurrency) so runs queue instead of overrunning the database. Direct connection limits for common Postgres providers: +Keep `concurrent runs × pool size` under your provider's connection limit, and cap how many runs execute at once with [concurrency limits](/concurrency) so runs queue instead of overrunning the database. Direct connection limits for common Postgres providers: | Provider | Direct connection limit | | --- | --- | @@ -197,7 +197,7 @@ export const myChat = chat.agent({ ## Troubleshooting -`too many connections` or connection refused: `concurrent runs × pool size` is over your provider's limit. Lower the pool size, cap [concurrency](/queue-concurrency), or connect through a pooler. +`too many connections` or connection refused: `concurrent runs × pool size` is over your provider's limit. Lower the pool size, cap [concurrency](/concurrency), or connect through a pooler. The worker crashes right after resuming from a wait: an idle connection that closed during the suspend emitted an unhandled `error` event. Attach `pool.on("error", ...)` on a `pg` pool (node-postgres or Drizzle); Prisma and the MongoDB driver handle this internally. @@ -208,6 +208,6 @@ When a task waits, the runtime can [checkpoint](/how-it-works#the-checkpoint-res ## See also - [Wait](/wait) for the primitives that trigger a checkpoint. -- [Concurrency and queues](/queue-concurrency) to cap how many runs execute at once. +- [Concurrency](/concurrency) to cap how many runs execute at once. - [Lifecycle functions](/tasks/overview#onwait-and-onresume-functions) for global `tasks.onWait` and `tasks.onResume`. - [Chat agent lifecycle hooks](/ai-chat/lifecycle-hooks) for `onChatSuspend` and `onChatResume`. diff --git a/docs/docs.json b/docs/docs.json index 13a867be930..6707e22a016 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -65,7 +65,8 @@ "group": "Wait", "pages": ["wait", "wait-for", "wait-until", "wait-for-token"] }, - "queue-concurrency", + "queues", + "concurrency", "versioning", "machines", "idempotency", @@ -389,6 +390,15 @@ "management/queues/concurrency-reset" ] }, + { + "group": "Concurrency limits API", + "pages": [ + "management/concurrency-limits/list", + "management/concurrency-limits/retrieve", + "management/concurrency-limits/override", + "management/concurrency-limits/reset" + ] + }, { "group": "Schedules API", "pages": [ @@ -737,6 +747,11 @@ ] }, "redirects": [ + { + "source": "/queue-concurrency", + "destination": "/concurrency", + "permanent": true + }, { "source": "/config/extensions", "destination": "/config/extensions/overview", diff --git a/docs/idempotency.mdx b/docs/idempotency.mdx index 0d61341697a..26065505ee2 100644 --- a/docs/idempotency.mdx +++ b/docs/idempotency.mdx @@ -3,7 +3,9 @@ title: "Idempotency" description: "An API call or operation is idempotent if it has the same result when called more than once." --- -We currently support idempotency at the task level, meaning that if you trigger a task with the same `idempotencyKey` twice, the second request will not create a new task run. Instead, the original run's handle is returned, allowing you to track the existing run's progress. +**Idempotency keys deduplicate task triggers.** If you trigger a task with the same `idempotencyKey` twice, the second request does not create a new run. It returns the original run's handle, so you can track the existing run's progress. + +The key applies to the trigger call and nothing else. It deduplicates `trigger()`, `triggerAndWait()`, and `batchTrigger()`; it does not make the code inside a task's `run()` function idempotent. See [Side effects inside `run()`](#side-effects-inside-run) for what that means when a task calls an external API. ## Why use idempotency keys? @@ -31,10 +33,46 @@ sequenceDiagram Other common use cases include: -- **Preventing duplicate emails** - Ensure a confirmation email is only sent once, even if the parent task retries -- **Avoiding double-charging customers** - Prevent duplicate payment processing during retries -- **One-time setup tasks** - Ensure initialization or migration tasks only run once -- **Deduplicating webhook processing** - Handle the same webhook event only once, even if it's delivered multiple times +- **Preventing duplicate emails** - Trigger the email task once, even if the parent task retries +- **Avoiding double-charging customers** - Trigger the payment task once, even if the parent task retries. The payment API call itself still needs the provider's own idempotency key, see [Side effects inside `run()`](#side-effects-inside-run) +- **One-time setup tasks** - Trigger initialization or migration tasks once +- **Deduplicating webhook processing** - Trigger the handler task once per webhook event, even if the event is delivered multiple times + +## Side effects inside `run()` + +An idempotency key stops a task from being *triggered* twice. It does not stop the code inside that task from *running* twice. When a run retries, `run()` executes again from the top, and every API call, database write, or email inside it happens again unless that operation is itself idempotent. + +This matters most for payments. A refund issued directly inside a retryable task is issued again on every retry, even when the retry is caused by an unrelated failure later in the same function. Moving the refund into a child task and triggering it with an idempotency key does not close the gap either: the child is triggered once, but its own `run()` can retry after the provider has accepted the request. Only the provider can deduplicate its own API call. + +Pass the provider's idempotency key on every call that must happen at most once. Derive it from a value that stays the same across every attempt of the run. The run ID, `ctx.run.id`, always qualifies: retries are attempts of the same run, so the ID does not change between them. This is the same value the default `run` scope mixes into a Trigger.dev idempotency key. + +```ts /trigger/refund-order.ts +import { idempotencyKeys, task } from "@trigger.dev/sdk"; +import Stripe from "stripe"; +import { sendRefundEmail } from "./send-refund-email"; + +const stripe = new Stripe(process.env.STRIPE_SECRET_KEY!); + +export const refundOrder = task({ + id: "refund-order", + retry: { maxAttempts: 3 }, + run: async (payload: { orderId: string; paymentIntentId: string; amount: number }, { ctx }) => { + // The run ID is stable across retries, so Stripe returns the original refund instead of creating a second one + await stripe.refunds.create( + { payment_intent: payload.paymentIntentId, amount: payload.amount }, + { idempotencyKey: `refund-${ctx.run.id}` } + ); + + // The Trigger.dev key is run-scoped by default, so a retry does not create a second email run + const idempotencyKey = await idempotencyKeys.create("refund-email"); + await sendRefundEmail.trigger({ orderId: payload.orderId }, { idempotencyKey }); + }, +}); +``` + +The run ID protects against retries of one run. If the same task can be triggered more than once for the same order, either trigger it with a Trigger.dev idempotency key so repeated triggers resolve to the same run and the same run ID, or key the provider call on a business ID from your payload instead, such as `refund-${payload.orderId}`, which deduplicates across separate runs too. + +Use a Trigger.dev idempotency key to stop a retry triggering a task twice, and the provider's idempotency key to stop a retry calling the provider twice. For a side effect with no idempotency support, record that you performed it in your own database and check that record before performing it again. ## `idempotencyKey` option diff --git a/docs/introduction.mdx b/docs/introduction.mdx index 1baddd6a1fa..2a89efbaab4 100644 --- a/docs/introduction.mdx +++ b/docs/introduction.mdx @@ -51,7 +51,7 @@ Build durable, multi-turn agents with the [chat agent](/ai-chat/overview): one l ## Scale and scheduling -Set how many runs of a task execute at once, globally or per tenant, with [queues and concurrency](/queue-concurrency). Run a task on a [cron schedule](/tasks/scheduled) with timezone support, choose the CPU and memory each task runs on with [machines](/machines), and control what happens when a task throws with [errors and retries](/errors-retrying). +Set how many runs of a task execute at once, globally or per tenant, with [concurrency](/concurrency) and [queues](/queues). Run a task on a [cron schedule](/tasks/scheduled) with timezone support, choose the CPU and memory each task runs on with [machines](/machines), and control what happens when a task throws with [errors and retries](/errors-retrying). ## Self-hosting diff --git a/docs/management/concurrency-limits/list.mdx b/docs/management/concurrency-limits/list.mdx new file mode 100644 index 00000000000..b0f3ca2eba4 --- /dev/null +++ b/docs/management/concurrency-limits/list.mdx @@ -0,0 +1,4 @@ +--- +title: "List Concurrency Limits" +openapi: "v3-openapi GET /api/v1/concurrency-limits" +--- diff --git a/docs/management/concurrency-limits/override.mdx b/docs/management/concurrency-limits/override.mdx new file mode 100644 index 00000000000..b99d0f07196 --- /dev/null +++ b/docs/management/concurrency-limits/override.mdx @@ -0,0 +1,4 @@ +--- +title: "Override Concurrency Limit" +openapi: "v3-openapi POST /api/v1/concurrency-limits/{name}/override" +--- diff --git a/docs/management/concurrency-limits/reset.mdx b/docs/management/concurrency-limits/reset.mdx new file mode 100644 index 00000000000..606b36f0033 --- /dev/null +++ b/docs/management/concurrency-limits/reset.mdx @@ -0,0 +1,4 @@ +--- +title: "Reset Concurrency Limit" +openapi: "v3-openapi POST /api/v1/concurrency-limits/{name}/reset" +--- diff --git a/docs/management/concurrency-limits/retrieve.mdx b/docs/management/concurrency-limits/retrieve.mdx new file mode 100644 index 00000000000..47b2e570866 --- /dev/null +++ b/docs/management/concurrency-limits/retrieve.mdx @@ -0,0 +1,4 @@ +--- +title: "Retrieve Concurrency Limit" +openapi: "v3-openapi GET /api/v1/concurrency-limits/{name}" +--- diff --git a/docs/queue-concurrency.mdx b/docs/queue-concurrency.mdx deleted file mode 100644 index b832ffc26da..00000000000 --- a/docs/queue-concurrency.mdx +++ /dev/null @@ -1,358 +0,0 @@ ---- -title: "Concurrency & Queues" -description: "Configure what you want to happen when there is more than one run at a time." ---- - -When you trigger a task, it isn't executed immediately. Instead, the task [run](/runs) is placed into a queue for execution. - -By default, each task gets its own queue and the concurrency is only limited by your environment concurrency limit. If you need more control (for example, to limit concurrency or share limits across multiple tasks), you can define a custom queue as described later. - -Controlling concurrency is useful when you have a task that can't be run concurrently, or when you want to limit the number of runs to avoid overloading a resource. - -It's important to note that only actively executing runs count towards concurrency limits. Runs that are delayed or waiting in a queue do not consume concurrency slots until they begin execution. - -## Default concurrency - -By default, all tasks have an unbounded concurrency limit, limited only by the overall concurrency limits of your environment. - - - Your environment has a base concurrency limit and a burstable limit (default burst factor of 2.0x - the base limit). Individual queues are limited by the base concurrency limit, not the burstable - limit. For example, if your base limit is 10, your environment can burst up to 20 concurrent runs, - but any single queue can have at most 10 concurrent runs. If you're a paying customer you can - request higher burst limits by [contacting us](https://www.trigger.dev/contact). - - -## Setting task concurrency - -You can set the concurrency limit for a task by setting the `concurrencyLimit` property on the task's queue. This limits the number of runs that can be executing at any one time: - -```ts /trigger/one-at-a-time.ts -// This task will only run one at a time -export const oneAtATime = task({ - id: "one-at-a-time", - queue: { - concurrencyLimit: 1, - }, - run: async (payload) => { - //... - }, -}); -``` - -This is useful if you need to control access to a shared resource, like a database or an API that has rate limits. - -## Sharing concurrency between tasks - -As well as putting queue settings directly on a task, you can define a queue and reuse it across multiple tasks. This allows you to share the same concurrency limit: - -```ts /trigger/queue.ts -export const myQueue = queue({ - name: "my-queue", - concurrencyLimit: 1, -}); - -export const task1 = task({ - id: "task-1", - queue: myQueue, - run: async (payload: { message: string }) => { - // ... - }, -}); - -export const task2 = task({ - id: "task-2", - queue: myQueue, - run: async (payload: { message: string }) => { - // ... - }, -}); -``` - -In this example, `task1` and `task2` share the same queue, so only one of them can run at a time. - -## Setting the queue when you trigger a run - -When you trigger a task you can override the default queue. This is really useful if you sometimes have high priority runs. - -The task and queue definition: - -```ts /trigger/override-concurrency.ts -const paidQueue = queue({ - name: "paid-users", - concurrencyLimit: 10, -}); - -export const generatePullRequest = task({ - id: "generate-pull-request", - queue: { - //normally when triggering this task it will be limited to 1 run at a time - concurrencyLimit: 1, - }, - run: async (payload) => { - //todo generate a PR using OpenAI - }, -}); -``` - -Triggering from your backend and overriding the queue: - -```ts app/api/push/route.ts -import { generatePullRequest } from "~/trigger/override-concurrency"; - -export async function POST(request: Request) { - const data = await request.json(); - - if (data.branch === "main") { - //trigger the task, with the paid users queue - const handle = await generatePullRequest.trigger(data, { - // Set the paid users queue - queue: "paid-users", - }); - - return Response.json(handle); - } else { - //triggered with the default queue (concurrency of 1) - const handle = await generatePullRequest.trigger(data); - return Response.json(handle); - } -} -``` - -## Concurrency keys and per-tenant queuing - -If you're building an application where you want to run tasks for your users, you might want a separate queue for each of your users (or orgs, projects, etc.). - -You can do this by using `concurrencyKey`. It creates a copy of the queue for each unique value of the key. - -Your backend code: - -```ts app/api/pr/route.ts -import { generatePullRequest } from "~/trigger/override-concurrency"; - -export async function POST(request: Request) { - const data = await request.json(); - - if (data.isFreeUser) { - //the "free-users" queue has a concurrency limit of 1 - const handle = await generatePullRequest.trigger(data, { - queue: "free-users", - //this creates a free-users queue for each user - concurrencyKey: data.userId, - }); - - //return a success response with the handle - return Response.json(handle); - } else { - //the "paid-users" queue has a concurrency limit of 10 - const handle = await generatePullRequest.trigger(data, { - queue: "paid-users", - //this creates a paid-users queue for each user - concurrencyKey: data.userId, - }); - - //return a success response with the handle - return Response.json(handle); - } -} -``` - -## Concurrency and subtasks - -When you trigger a task that has subtasks, the subtasks will not inherit the queue from the parent task. Unless otherwise specified, subtasks will run on their own queue - -```ts /trigger/subtasks.ts -export const parentTask = task({ - id: "parent-task", - run: async (payload) => { - //trigger a subtask - await subtask.triggerAndWait(payload); - }, -}); - -// This subtask will run on its own queue -export const subtask = task({ - id: "subtask", - run: async (payload) => { - //... - }, -}); -``` - -## Waits and concurrency - -With our [task checkpoint system](/how-it-works#the-checkpoint-resume-system), tasks can wait at various waitpoints (like waiting for subtasks to complete, delays, or external events). The way this system interacts with the concurrency system is important to understand. - -Concurrency is only released when a run reaches a waitpoint and is checkpointed. When a run is checkpointed, it transitions to the `WAITING` state and releases its concurrency slot back to both the queue and the environment, allowing other runs to execute or resume. - -This means that: - -- Only actively executing runs count towards concurrency limits -- Runs in the `WAITING` state (checkpointed at waitpoints) do not consume concurrency slots -- You can have more runs in the `WAITING` state than your queue's concurrency limit -- When a waiting run resumes (e.g., when a subtask completes), it must re-acquire a concurrency slot - -For example, if you have a queue with a `concurrencyLimit` of 1: - -- You can only have exactly 1 run executing at a time -- You may have multiple runs in the `WAITING` state that belong to that queue -- When the executing run reaches a waitpoint and checkpoints, it releases its slot -- The next queued run can then begin execution - -### Short time-based waits keep their slot - -Checkpointing takes time, so a run doesn't checkpoint the moment it reaches a waitpoint. For [`wait.for()`](/wait-for) and [`wait.until()`](/wait-until) it happens 60 seconds into the wait, so anything shorter stays `EXECUTING` and holds its slot for the whole wait. If you're polling in a loop, use an interval comfortably above 60 seconds so the slot is actually released between polls. - -### Waiting for a subtask on a different queue - -When a parent task triggers and waits for a subtask on a different queue, the parent task will checkpoint and release its concurrency slot once it reaches the wait point. This prevents environment deadlocks where all concurrency slots would be occupied by waiting tasks. - -```ts /trigger/waiting.ts -export const parentTask = task({ - id: "parent-task", - queue: { - concurrencyLimit: 1, - }, - run: async (payload) => { - //trigger a subtask and wait for it to complete - await subtask.triggerAndWait(payload); - // The parent task checkpoints here and releases its concurrency slot - // allowing other tasks to execute while waiting - }, -}); - -export const subtask = task({ - id: "subtask", - run: async (payload) => { - //... - }, -}); -``` - -When the parent task reaches the `triggerAndWait` call, it checkpoints and transitions to the `WAITING` state, releasing its concurrency slot back to both its queue and the environment. Once the subtask completes, the parent task will resume and re-acquire a concurrency slot. - -## Managing queues with the SDK - -The SDK provides a `queues` namespace that allows you to manage queues programmatically. You can list, retrieve, pause, resume, and modify concurrency limits for queues. - - - Import from `@trigger.dev/sdk`: - ```ts - import { queues } from "@trigger.dev/sdk"; - ``` - - -### Listing queues - -You can list all queues in your environment with pagination support: - -```ts -import { queues } from "@trigger.dev/sdk"; - -// List all queues (returns paginated results) -const allQueues = await queues.list(); - -// With pagination options -const pagedQueues = await queues.list({ - page: 1, - perPage: 20, -}); -``` - -### Retrieving a queue - -You can retrieve a specific queue by its ID, or by its type and name: - -```ts -import { queues } from "@trigger.dev/sdk"; - -// Using queue ID (starts with "queue_") -const queueById = await queues.retrieve("queue_1234"); - -// Using type and name for a task's default queue -const taskQueue = await queues.retrieve({ - type: "task", - name: "my-task-id", -}); - -// Using type and name for a custom queue -const customQueue = await queues.retrieve({ - type: "custom", - name: "my-custom-queue", -}); -``` - -The queue object contains useful information about the queue state: - -```ts -{ - id: "queue_1234", // Queue ID - name: "my-task-id", // Queue name - type: "task", // "task" or "custom" - running: 5, // Currently executing runs - queued: 10, // Runs waiting to execute - paused: false, // Whether the queue is paused - concurrencyLimit: 10, // Current concurrency limit - concurrency: { - current: 10, // Effective limit - base: 10, // Default limit from code - override: null, // Override value (if set) - overriddenAt: null, // When override was applied - overriddenBy: null, // Who applied the override - } -} -``` - -### Pausing and resuming queues - -You can pause a queue to prevent new runs from starting. Runs that are currently executing will continue to completion. - -```ts -import { queues } from "@trigger.dev/sdk"; - -// Pause a queue using its ID -await queues.pause("queue_1234"); - -// Or using type and name -await queues.pause({ type: "task", name: "my-task-id" }); -await queues.pause({ type: "custom", name: "my-custom-queue" }); -``` - -To resume a paused queue and allow new runs to start: - -```ts -import { queues } from "@trigger.dev/sdk"; - -// Resume a queue using its ID -await queues.resume("queue_1234"); - -// Or using type and name -await queues.resume({ type: "task", name: "my-task-id" }); -await queues.resume({ type: "custom", name: "my-custom-queue" }); -``` - -### Overriding concurrency limits - -You can temporarily override a queue's concurrency limit. This is useful for scaling up or down based on demand: - -```ts -import { queues } from "@trigger.dev/sdk"; - -// Set concurrency limit to 5 -await queues.overrideConcurrencyLimit("queue_1234", 5); - -// Or using type and name -await queues.overrideConcurrencyLimit({ type: "task", name: "my-task-id" }, 20); -``` - -To reset the concurrency limit back to the base value defined in your code: - -```ts -import { queues } from "@trigger.dev/sdk"; - -// Reset concurrency limit to the base value -await queues.resetConcurrencyLimit("queue_1234"); - -// Or using type and name -await queues.resetConcurrencyLimit({ type: "task", name: "my-task-id" }); -``` diff --git a/docs/queues.mdx b/docs/queues.mdx new file mode 100644 index 00000000000..7e50d589c91 --- /dev/null +++ b/docs/queues.mdx @@ -0,0 +1,237 @@ +--- +title: "Queues" +description: "Control the order your runs execute in." +--- + +import Priority from "/snippets/priority.mdx"; + +When you trigger a task, it isn't executed immediately. Instead, the task [run](/runs) is placed into a queue for execution. Runs in a queue are started in the order they were triggered, first in, first out, unless you give a run a [priority](#priority) at trigger time. + +By default, each task gets its own queue, so triggering the same task ten times executes those runs in trigger order. How many of them execute *at once* is a separate question, governed by [concurrency](/concurrency) — your environment's concurrency limit, plus any limits you set on the task. + +## Sharing a queue between tasks + +Declare a queue with `queue()` and set it on several tasks to interleave their runs in one queue. Runs from every task on the queue start in the order they were triggered, whichever task they belong to: + +```ts /trigger/emails.ts +import { queue, task } from "@trigger.dev/sdk"; + +const emailQueue = queue({ name: "emails" }); + +export const sendWelcomeEmail = task({ + id: "send-welcome-email", + queue: emailQueue, + run: async (payload) => { + //... + }, +}); + +export const sendDigestEmail = task({ + id: "send-digest-email", + queue: emailQueue, + run: async (payload) => { + //... + }, +}); +``` + + + If what you want to share is a concurrency cap rather than ordering, you don't need a shared + queue: declare a named limit with `concurrencyLimit()` and each task keeps its own queue while + drawing from the shared limit. See [sharing a limit between + tasks](/concurrency#sharing-a-limit-between-tasks). + + +Queue names can use letters, numbers, underscores, hyphens and slashes, up to 128 characters. Names starting with `limit/` are reserved for [concurrency limits](/concurrency) and are rejected at deploy time. + +## Setting the queue when you trigger a run + +When you trigger a task you can override its queue by name. This is really useful if you sometimes have high priority runs: + +```ts /trigger/override-queue.ts +import { queue, task } from "@trigger.dev/sdk"; + +const paidQueue = queue({ name: "paid-users" }); + +export const generatePullRequest = task({ + id: "generate-pull-request", + // normally this task is limited to 1 run at a time + concurrency: { total: 1 }, + run: async (payload) => { + //todo generate a PR using OpenAI + }, +}); +``` + +Triggering from your backend and overriding the queue: + +```ts app/api/push/route.ts +import { generatePullRequest } from "~/trigger/override-queue"; + +export async function POST(request: Request) { + const data = await request.json(); + + if (data.branch === "main") { + //trigger the task, with the paid users queue + const handle = await generatePullRequest.trigger(data, { + // Set the paid users queue + queue: "paid-users", + }); + + return Response.json(handle); + } else { + //triggered with the default queue + const handle = await generatePullRequest.trigger(data); + return Response.json(handle); + } +} +``` + +## Priority + +You can re-order runs within a queue at trigger time by giving them a priority: + + + +## Managing queues with the SDK + +The SDK provides a `queues` namespace that allows you to manage queues programmatically. You can list, retrieve, pause, resume, and modify concurrency limits for queues. + + + Import from `@trigger.dev/sdk`: + ```ts + import { queues } from "@trigger.dev/sdk"; + ``` + + +### Listing queues + +You can list all queues in your environment with pagination support: + +```ts +import { queues } from "@trigger.dev/sdk"; + +// List all queues (returns paginated results) +const allQueues = await queues.list(); + +// With pagination options +const pagedQueues = await queues.list({ + page: 1, + perPage: 20, +}); +``` + +### Retrieving a queue + +You can retrieve a specific queue by its ID, or by its type and name: + +```ts +import { queues } from "@trigger.dev/sdk"; + +// Using queue ID (starts with "queue_") +const queueById = await queues.retrieve("queue_1234"); + +// Using type and name for a task's default queue +const taskQueue = await queues.retrieve({ + type: "task", + name: "my-task-id", +}); + +// Using type and name for a custom queue +const customQueue = await queues.retrieve({ + type: "custom", + name: "my-custom-queue", +}); +``` + +The queue object contains useful information about the queue state, and its `version` discriminates the shape: + +```ts +// V1: the queue carries its own concurrency limit and override state +{ + id: "queue_1234", // Queue ID + name: "my-task-id", // Queue name + type: "task", // "task" or "custom" + version: "V1", + running: 5, // Currently executing runs + queued: 10, // Runs waiting to execute + paused: false, // Whether the queue is paused + concurrencyLimit: 10, // The queue's own limit + concurrency: { + current: 10, // Effective limit + base: 10, // Default limit from code + override: null, // Override value (if set) + overriddenAt: null, // When override was applied + overriddenBy: null, // Who applied the override + } +} + +// V2: the queue is only the line runs wait in. Concurrency is declared with +// the task `concurrency` option and managed through `concurrencyLimits`. +{ + id: "queue_5678", + name: "my-v2-task-id", + type: "task", + version: "V2", + running: 3, + queued: 12, + paused: false, + concurrencyLimit: null, +} +``` + +### Pausing and resuming queues + +You can pause a queue to prevent new runs from starting. Runs that are currently executing will continue to completion. + +```ts +import { queues } from "@trigger.dev/sdk"; + +// Pause a queue using its ID +await queues.pause("queue_1234"); + +// Or using type and name +await queues.pause({ type: "task", name: "my-task-id" }); +await queues.pause({ type: "custom", name: "my-custom-queue" }); +``` + +To resume a paused queue and allow new runs to start: + +```ts +import { queues } from "@trigger.dev/sdk"; + +// Resume a queue using its ID +await queues.resume("queue_1234"); + +// Or using type and name +await queues.resume({ type: "task", name: "my-task-id" }); +await queues.resume({ type: "custom", name: "my-custom-queue" }); +``` + +### Overriding concurrency limits + +You can temporarily override a queue's concurrency limit. This is useful for scaling up or down based on demand: + +```ts +import { queues } from "@trigger.dev/sdk"; + +// Set concurrency limit to 5 +await queues.overrideConcurrencyLimit("queue_1234", 5); + +// Or using type and name +await queues.overrideConcurrencyLimit({ type: "task", name: "my-task-id" }, 20); +``` + +To reset the concurrency limit back to the base value defined in your code: + +```ts +import { queues } from "@trigger.dev/sdk"; + +// Reset concurrency limit to the base value +await queues.resetConcurrencyLimit("queue_1234"); + +// Or using type and name +await queues.resetConcurrencyLimit({ type: "task", name: "my-task-id" }); +``` + +Overrides survive deploys: redeploying your code keeps an active override until you reset it. diff --git a/docs/reports.mdx b/docs/reports.mdx index 07d517bc3b0..c4a330a569c 100644 --- a/docs/reports.mdx +++ b/docs/reports.mdx @@ -148,7 +148,7 @@ An unknown report key returns `404` with the list of available keys. Every tool the MCP server exposes, including `get_report`. - + Configure the concurrency limits the Flow verdict checks against. diff --git a/docs/runs/priority.mdx b/docs/runs/priority.mdx index 20da1894629..374b248078f 100644 --- a/docs/runs/priority.mdx +++ b/docs/runs/priority.mdx @@ -3,29 +3,6 @@ title: "Priority" description: "Specify a priority when triggering a run." --- -You can set a priority when you trigger a run. This allows you to prioritize some of your runs over others, so they are started sooner. This is very useful when: +import Priority from "/snippets/priority.mdx"; -- You have critical work that needs to start more quickly (and you have long queues). -- You want runs for your premium users to take priority over free users. - -The value for priority is a time offset in seconds that determines the order of dequeuing. - -![Priority runs](/images/priority-runs.png) - -If you specify a priority of `10` the run will dequeue before runs that were triggered with no priority 8 seconds ago, like in this example: - -```ts -// no priority = 0 -await myTask.trigger({ foo: "bar" }); - -//... imagine 8s pass by - -// this run will start before the run above that was triggered 8s ago (with no priority) -await myTask.trigger({ foo: "bar" }, { priority: 10 }); -``` - -If you passed a value of `3600` the run would dequeue before runs that were triggered an hour ago (with no priority). - - - Setting a high priority will not allow you to beat runs from other organizations. It will only affect the order of your own runs. - + diff --git a/docs/snippets/priority.mdx b/docs/snippets/priority.mdx new file mode 100644 index 00000000000..91d99a6de6e --- /dev/null +++ b/docs/snippets/priority.mdx @@ -0,0 +1,27 @@ +You can set a priority when you trigger a run. This allows you to prioritize some of your runs over others, so they are started sooner. This is very useful when: + +- You have critical work that needs to start more quickly (and you have long queues). +- You want runs for your premium users to take priority over free users. + +The value for priority is a time offset in seconds that determines the order of dequeuing. + +![Priority runs](/images/priority-runs.png) + +If you specify a priority of `10` the run will dequeue before runs that were triggered with no priority 8 seconds ago, like in this example: + +```ts +// no priority = 0 +await myTask.trigger({ foo: "bar" }); + +//... imagine 8s pass by + +// this run will start before the run above that was triggered 8s ago (with no priority) +await myTask.trigger({ foo: "bar" }, { priority: 10 }); +``` + +If you passed a value of `3600` the run would dequeue before runs that were triggered an hour ago (with no priority). + + + Setting a high priority will not allow you to beat runs from other organizations. It will only + affect the order of your own runs. + diff --git a/docs/tasks/overview.mdx b/docs/tasks/overview.mdx index d7cb4d8f4d9..9d287a93334 100644 --- a/docs/tasks/overview.mdx +++ b/docs/tasks/overview.mdx @@ -98,16 +98,14 @@ For more information read [the retrying guide](/errors-retrying). It's also worth mentioning that you can [retry a block of code](/errors-retrying) inside your tasks as well. -### `queue` options +### `queue` and `concurrency` options -Queues allow you to control the concurrency of your tasks. This allows you to have one-at-a-time execution and parallel executions. There are also more advanced techniques like having different concurrencies for different sets of your users. For more information read [the concurrency & queues guide](/queue-concurrency). +Queues control the order your runs execute in, and concurrency limits control how many execute at once: one-at-a-time execution, parallel executions, and more advanced techniques like separate concurrency for different sets of your users. For more information read the [Queues](/queues) and [Concurrency](/concurrency) guides. ```ts /trigger/one-at-a-time.ts export const oneAtATime = task({ id: "one-at-a-time", - queue: { - concurrencyLimit: 1, - }, + concurrency: { total: 1 }, run: async (payload: any, { ctx }) => { //... }, diff --git a/docs/v3-openapi.yaml b/docs/v3-openapi.yaml index 6c182fad5bc..8bc07bd64bd 100644 --- a/docs/v3-openapi.yaml +++ b/docs/v3-openapi.yaml @@ -3074,7 +3074,8 @@ paths: description: The queue ID (e.g., `queue_1234`), or the name of the queue when using the `type` body parameter. example: queue_1234 requestBody: - required: false + required: true + description: At least an empty JSON object `{}` must be sent; a zero-length body is rejected with a 400. content: application/json: schema: @@ -3120,6 +3121,203 @@ paths: name: "my-task-id", }); + "/api/v1/concurrency-limits": + get: + operationId: list_concurrency_limits_v1 + summary: List concurrency limits + description: | + List the environment's declared concurrency limits (anonymous inline limits + appear under their derived `task/` names), with each limit's bounds + and its live running and queued counts. Results are ordered by the underlying + row name, so named limits sort before `task/`-derived inline limits. + parameters: + - name: page + in: query + required: false + schema: + type: integer + minimum: 1 + default: 1 + - name: perPage + in: query + required: false + schema: + type: integer + minimum: 1 + default: 25 + maximum: 100 + responses: + "200": + description: Concurrency limits retrieved successfully + content: + application/json: + schema: + type: object + properties: + data: + type: array + items: + "$ref": "#/components/schemas/ConcurrencyLimitObject" + pagination: + type: object + properties: + currentPage: + type: integer + totalPages: + type: integer + count: + type: integer + "401": + description: Unauthorized request + tags: + - concurrency-limits + security: + - secretKey: [] + x-codeSamples: + - lang: typescript + source: |- + import { concurrencyLimits } from "@trigger.dev/sdk"; + + const limits = await concurrencyLimits.list({ page: 1, perPage: 20 }); + + "/api/v1/concurrency-limits/{name}": + get: + operationId: retrieve_concurrency_limit_v1 + summary: Retrieve concurrency limit + description: Retrieve a concurrency limit by name, with its bounds and live running and queued counts. + parameters: + - name: name + in: path + required: true + schema: + type: string + description: The limit's name, as declared with `concurrencyLimit()` + responses: + "200": + description: Concurrency limit retrieved successfully + content: + application/json: + schema: + "$ref": "#/components/schemas/ConcurrencyLimitObject" + "401": + description: Unauthorized request + "404": + description: Concurrency limit not found + tags: + - concurrency-limits + security: + - secretKey: [] + x-codeSamples: + - lang: typescript + source: |- + import { concurrencyLimits } from "@trigger.dev/sdk"; + + const limit = await concurrencyLimits.retrieve("openai"); + + "/api/v1/concurrency-limits/{name}/override": + post: + operationId: override_concurrency_limit_v1 + summary: Override concurrency limit + description: | + Override a concurrency limit's bounds. Only the given fields change; the declared + values are kept and restored by reset. Overriding `total` to `0` blocks every run + holding the limit, which is how a limit is paused. + parameters: + - name: name + in: path + required: true + schema: + type: string + description: The limit's name + requestBody: + required: true + content: + application/json: + schema: + type: object + minProperties: 1 + additionalProperties: false + description: At least one of `perKey` or `total` must be provided. + properties: + perKey: + type: integer + minimum: 0 + maximum: 100000 + description: Caps each concurrencyKey pool; runs without a key share one pool. May not exceed the environment concurrency limit. + total: + type: integer + minimum: 0 + maximum: 100000 + description: Caps every run holding this limit, keys or not. May not exceed the environment concurrency limit. + responses: + "200": + description: Concurrency limit overridden successfully + content: + application/json: + schema: + "$ref": "#/components/schemas/ConcurrencyLimitObject" + "400": + description: Invalid request parameters, or a bound exceeds the environment concurrency limit + "401": + description: Unauthorized request + "404": + description: Concurrency limit not found + "409": + description: The limit changed concurrently; retry the request + tags: + - concurrency-limits + security: + - secretKey: [] + x-codeSamples: + - lang: typescript + source: |- + import { concurrencyLimits } from "@trigger.dev/sdk"; + + // Raise the shared cap + await concurrencyLimits.override("openai", { total: 50 }); + + // Pause the limit + await concurrencyLimits.override("openai", { total: 0 }); + + "/api/v1/concurrency-limits/{name}/reset": + post: + operationId: reset_concurrency_limit_v1 + summary: Reset concurrency limit + description: Reset a concurrency limit back to the values declared in your code, clearing any override. + parameters: + - name: name + in: path + required: true + schema: + type: string + description: The limit's name + responses: + "200": + description: Concurrency limit reset successfully + content: + application/json: + schema: + "$ref": "#/components/schemas/ConcurrencyLimitObject" + "400": + description: The limit has no override to reset + "401": + description: Unauthorized request + "404": + description: Concurrency limit not found + "409": + description: The limit changed concurrently; retry the request + tags: + - concurrency-limits + security: + - secretKey: [] + x-codeSamples: + - lang: typescript + source: |- + import { concurrencyLimits } from "@trigger.dev/sdk"; + + await concurrencyLimits.reset("openai"); + + "/api/v1/waitpoints/tokens": post: operationId: create_waitpoint_token_v1 @@ -4275,9 +4473,85 @@ components: minimum: 0 maximum: 1000 description: 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. + ConcurrencyLimitObject: + type: object + required: + - id + - name + - perKey + - total + - running + - queued + properties: + id: + type: string + description: The limit's id, starting with `climit_` + example: "climit_abcdef123456" + name: + type: string + description: The limit's name, as declared with `concurrencyLimit()` + example: "openai" + perKey: + type: object + description: Caps each concurrencyKey pool; runs without a key share one pool + properties: + current: + type: integer + nullable: true + description: Enforced right now (null = no per-key bound). Enforcement clamps it to the environment concurrency limit at admit time. + example: null + base: + type: integer + nullable: true + description: The declared value an override reverts to on reset + example: null + override: + type: integer + nullable: true + description: The overridden value, when an override is active + example: null + overriddenAt: + type: string + format: date-time + nullable: true + description: When the override was applied + example: null + total: + type: object + description: Caps every run holding this limit, keys or not + properties: + current: + type: integer + nullable: true + description: Enforced right now (null = no total bound). Enforcement clamps it to the environment concurrency limit at admit time. + example: 25 + base: + type: integer + nullable: true + description: The declared value an override reverts to on reset + example: 25 + override: + type: integer + nullable: true + description: The overridden value, when an override is active + example: null + overriddenAt: + type: string + format: date-time + nullable: true + description: When the override was applied + example: null + running: + type: integer + description: Runs executing that hold this limit + example: 14 + queued: + type: integer + description: Runs that are queued and must clear this limit to execute + example: 100 QueueObject: type: object - required: ["id", "name", "type", "running", "queued", "paused"] + required: ["id", "name", "type", "version", "running", "queued", "paused"] properties: id: type: string @@ -4295,6 +4569,14 @@ components: - `task`: Created automatically for each task - `custom`: Created explicitly in your code using `queue()` example: task + version: + type: string + enum: [V1, V2] + description: | + Discriminates the shape: + - `V1`: the queue carries its own `concurrencyLimit` (applied per key when runs pass a `concurrencyKey`, to the whole queue when they don't) and its `concurrency` override state + - `V2`: the queue is only the line runs wait in; concurrency is declared with the task `concurrency` option and managed through the concurrency-limits endpoints. `concurrencyLimit` is always null and `concurrency` is absent. + example: V1 running: type: integer description: The number of runs currently executing @@ -4310,11 +4592,11 @@ components: concurrencyLimit: type: integer nullable: true - description: The current concurrency limit of the queue + description: The queue's own concurrency limit. Meaningful on V1 queues; always null on V2 queues. example: 10 concurrency: type: object - description: Detailed concurrency information + description: Detailed concurrency information. V1 queues only. properties: current: type: integer diff --git a/docs/writing-tasks-introduction.mdx b/docs/writing-tasks-introduction.mdx index 5f0bc330912..fe738d4db23 100644 --- a/docs/writing-tasks-introduction.mdx +++ b/docs/writing-tasks-introduction.mdx @@ -15,7 +15,8 @@ Before digging deeper into the details of writing tasks, you should read the [fu | [Logging](/logging) | View and send logs and traces from your tasks. | | [Errors & retrying](/errors-retrying) | How to deal with errors and write reliable tasks. | | [Wait](/wait) | Wait for periods of time or for external events to occur before continuing. | -| [Concurrency & Queues](/queue-concurrency) | Configure what you want to happen when there is more than one run at a time. | +| [Queues](/queues) | Control the order your runs execute in. | +| [Concurrency](/concurrency) | Limit how many runs execute at once: per task, per tenant, or shared across tasks. | | [Realtime notifications](/realtime/overview) | Send realtime notifications from your task that you can subscribe to from your backend or frontend. | | [Versioning](/versioning) | How versioning works. | | [Machines](/machines) | Configure the CPU and RAM of the machine your task runs on | diff --git a/internal-packages/clickhouse/schema/048_add_queue_metrics_total_concurrency.sql b/internal-packages/clickhouse/schema/048_add_queue_metrics_total_concurrency.sql new file mode 100644 index 00000000000..6a7abc8d9b9 --- /dev/null +++ b/internal-packages/clickhouse/schema/048_add_queue_metrics_total_concurrency.sql @@ -0,0 +1,178 @@ +-- +goose Up + +-- Total-concurrency gauges: total_running is the in-flight count across ALL +-- concurrency-key variants of a queue (the groupConcurrency set), total_limit the +-- RAW stored total cap (0 = none, readers clamp against max_env_limit). Emitted on +-- base-queue gauge rows only. Per-key gauge rows carry the queue concurrency +-- limit that applied in queue_limit, surfaced in the ck tier as max_limit +-- (1000000 = no explicit limit). + +ALTER TABLE trigger_dev.queue_metrics_raw_v1 + ADD COLUMN IF NOT EXISTS total_running UInt32 DEFAULT 0, + ADD COLUMN IF NOT EXISTS total_limit UInt32 DEFAULT 0; + +ALTER TABLE trigger_dev.queue_metrics_v1 + ADD COLUMN IF NOT EXISTS max_total_running SimpleAggregateFunction(max, UInt32), + ADD COLUMN IF NOT EXISTS max_total_limit SimpleAggregateFunction(max, UInt32); + +ALTER TABLE trigger_dev.queue_metrics_5m_v1 + ADD COLUMN IF NOT EXISTS max_total_running SimpleAggregateFunction(max, UInt32), + ADD COLUMN IF NOT EXISTS max_total_limit SimpleAggregateFunction(max, UInt32); + +ALTER TABLE trigger_dev.queue_metrics_ck_v1 + ADD COLUMN IF NOT EXISTS max_limit SimpleAggregateFunction(max, UInt32); + +-- Materialized views cannot be altered: recreate them with the new columns. The 5m +-- MV MUST keep reading raw, never cascade off queue_metrics_v1 (out-of-time-order +-- deltaSumTimestamp merges double-count bridging spans). + +DROP VIEW IF EXISTS trigger_dev.queue_metrics_mv_v1; +CREATE MATERIALIZED VIEW IF NOT EXISTS trigger_dev.queue_metrics_mv_v1 +TO trigger_dev.queue_metrics_v1 AS +SELECT + organization_id, project_id, environment_id, queue_name, + toStartOfInterval(event_time, INTERVAL 10 SECOND) AS bucket_start, + deltaSumTimestampStateIf(cumulative, order_key, op = 'enqueue' AND concurrency_key = '') AS enqueue_delta, + deltaSumTimestampStateIf(cumulative, order_key, op = 'started' AND concurrency_key = '') AS started_delta, + deltaSumTimestampStateIf(cumulative, order_key, op = 'ack' AND concurrency_key = '') AS ack_delta, + deltaSumTimestampStateIf(cumulative, order_key, op = 'nack' AND concurrency_key = '') AS nack_delta, + deltaSumTimestampStateIf(cumulative, order_key, op = 'dlq' AND concurrency_key = '') AS dlq_delta, + sum(throttled) AS throttled_count, + max(queued) AS max_queued, + max(running) AS max_running, + max(queue_limit) AS max_limit, + max(env_queued) AS max_env_queued, + max(env_running) AS max_env_running, + max(env_limit) AS max_env_limit, + max(ck_backlogged) AS max_ck_backlogged, + max(ck_max_wait_ms) AS max_ck_wait_ms, + max(total_running) AS max_total_running, + max(total_limit) AS max_total_limit, + sumIf(wait_ms, op = 'started' AND concurrency_key = '') AS wait_ms_sum, + countIf(op = 'started' AND wait_ms > 0 AND concurrency_key = '') AS wait_ms_count, + quantilesStateIf(0.5, 0.9, 0.95, 0.99)(wait_ms, op = 'started' AND wait_ms > 0 AND concurrency_key = '') AS wait_quantiles +FROM trigger_dev.queue_metrics_raw_v1 +GROUP BY organization_id, project_id, environment_id, queue_name, bucket_start; + +DROP VIEW IF EXISTS trigger_dev.queue_metrics_5m_mv_v1; +CREATE MATERIALIZED VIEW IF NOT EXISTS trigger_dev.queue_metrics_5m_mv_v1 +TO trigger_dev.queue_metrics_5m_v1 AS +SELECT + organization_id, project_id, environment_id, queue_name, + toStartOfInterval(event_time, INTERVAL 5 MINUTE) AS bucket_start, + deltaSumTimestampStateIf(cumulative, order_key, op = 'enqueue' AND concurrency_key = '') AS enqueue_delta, + deltaSumTimestampStateIf(cumulative, order_key, op = 'started' AND concurrency_key = '') AS started_delta, + deltaSumTimestampStateIf(cumulative, order_key, op = 'ack' AND concurrency_key = '') AS ack_delta, + deltaSumTimestampStateIf(cumulative, order_key, op = 'nack' AND concurrency_key = '') AS nack_delta, + deltaSumTimestampStateIf(cumulative, order_key, op = 'dlq' AND concurrency_key = '') AS dlq_delta, + sum(throttled) AS throttled_count, + max(queued) AS max_queued, + max(running) AS max_running, + max(queue_limit) AS max_limit, + max(env_queued) AS max_env_queued, + max(env_running) AS max_env_running, + max(env_limit) AS max_env_limit, + max(ck_backlogged) AS max_ck_backlogged, + max(ck_max_wait_ms) AS max_ck_wait_ms, + max(total_running) AS max_total_running, + max(total_limit) AS max_total_limit, + sumIf(wait_ms, op = 'started' AND concurrency_key = '') AS wait_ms_sum, + countIf(op = 'started' AND wait_ms > 0 AND concurrency_key = '') AS wait_ms_count, + quantilesStateIf(0.5, 0.9, 0.95, 0.99)(wait_ms, op = 'started' AND wait_ms > 0 AND concurrency_key = '') AS wait_quantiles +FROM trigger_dev.queue_metrics_raw_v1 +GROUP BY organization_id, project_id, environment_id, queue_name, bucket_start; + +DROP VIEW IF EXISTS trigger_dev.queue_metrics_ck_mv_v1; +CREATE MATERIALIZED VIEW IF NOT EXISTS trigger_dev.queue_metrics_ck_mv_v1 +TO trigger_dev.queue_metrics_ck_v1 AS +SELECT + organization_id, project_id, environment_id, queue_name, concurrency_key, + toStartOfInterval(event_time, INTERVAL 10 SECOND) AS bucket_start, + deltaSumTimestampStateIf(cumulative, order_key, op = 'enqueue') AS enqueue_delta, + deltaSumTimestampStateIf(cumulative, order_key, op = 'started') AS started_delta, + deltaSumTimestampStateIf(cumulative, order_key, op = 'ack') AS ack_delta, + maxIf(queued, op = 'gauge') AS max_queued, + maxIf(running, op = 'gauge') AS max_running, + maxIf(queue_limit, op = 'gauge') AS max_limit, + sumIf(wait_ms, op = 'started') AS wait_ms_sum, + countIf(op = 'started' AND wait_ms > 0) AS wait_ms_count +FROM trigger_dev.queue_metrics_raw_v1 +WHERE concurrency_key != '' +GROUP BY organization_id, project_id, environment_id, queue_name, concurrency_key, bucket_start; + +-- +goose Down +DROP VIEW IF EXISTS trigger_dev.queue_metrics_ck_mv_v1; +DROP VIEW IF EXISTS trigger_dev.queue_metrics_5m_mv_v1; +DROP VIEW IF EXISTS trigger_dev.queue_metrics_mv_v1; +ALTER TABLE trigger_dev.queue_metrics_ck_v1 DROP COLUMN IF EXISTS max_limit; +ALTER TABLE trigger_dev.queue_metrics_5m_v1 DROP COLUMN IF EXISTS max_total_running, DROP COLUMN IF EXISTS max_total_limit; +ALTER TABLE trigger_dev.queue_metrics_v1 DROP COLUMN IF EXISTS max_total_running, DROP COLUMN IF EXISTS max_total_limit; +ALTER TABLE trigger_dev.queue_metrics_raw_v1 DROP COLUMN IF EXISTS total_running, DROP COLUMN IF EXISTS total_limit; + +-- Recreate the pre-048 materialized views (the definitions from 036) so ingestion keeps +-- feeding every aggregate table after a rollback. +CREATE MATERIALIZED VIEW IF NOT EXISTS trigger_dev.queue_metrics_mv_v1 +TO trigger_dev.queue_metrics_v1 AS +SELECT + organization_id, project_id, environment_id, queue_name, + toStartOfInterval(event_time, INTERVAL 10 SECOND) AS bucket_start, + deltaSumTimestampStateIf(cumulative, order_key, op = 'enqueue' AND concurrency_key = '') AS enqueue_delta, + deltaSumTimestampStateIf(cumulative, order_key, op = 'started' AND concurrency_key = '') AS started_delta, + deltaSumTimestampStateIf(cumulative, order_key, op = 'ack' AND concurrency_key = '') AS ack_delta, + deltaSumTimestampStateIf(cumulative, order_key, op = 'nack' AND concurrency_key = '') AS nack_delta, + deltaSumTimestampStateIf(cumulative, order_key, op = 'dlq' AND concurrency_key = '') AS dlq_delta, + sum(throttled) AS throttled_count, + max(queued) AS max_queued, + max(running) AS max_running, + max(queue_limit) AS max_limit, + max(env_queued) AS max_env_queued, + max(env_running) AS max_env_running, + max(env_limit) AS max_env_limit, + max(ck_backlogged) AS max_ck_backlogged, + max(ck_max_wait_ms) AS max_ck_wait_ms, + sumIf(wait_ms, op = 'started' AND concurrency_key = '') AS wait_ms_sum, + countIf(op = 'started' AND wait_ms > 0 AND concurrency_key = '') AS wait_ms_count, + quantilesStateIf(0.5, 0.9, 0.95, 0.99)(wait_ms, op = 'started' AND wait_ms > 0 AND concurrency_key = '') AS wait_quantiles +FROM trigger_dev.queue_metrics_raw_v1 +GROUP BY organization_id, project_id, environment_id, queue_name, bucket_start; + +CREATE MATERIALIZED VIEW IF NOT EXISTS trigger_dev.queue_metrics_5m_mv_v1 +TO trigger_dev.queue_metrics_5m_v1 AS +SELECT + organization_id, project_id, environment_id, queue_name, + toStartOfInterval(event_time, INTERVAL 5 MINUTE) AS bucket_start, + deltaSumTimestampStateIf(cumulative, order_key, op = 'enqueue' AND concurrency_key = '') AS enqueue_delta, + deltaSumTimestampStateIf(cumulative, order_key, op = 'started' AND concurrency_key = '') AS started_delta, + deltaSumTimestampStateIf(cumulative, order_key, op = 'ack' AND concurrency_key = '') AS ack_delta, + deltaSumTimestampStateIf(cumulative, order_key, op = 'nack' AND concurrency_key = '') AS nack_delta, + deltaSumTimestampStateIf(cumulative, order_key, op = 'dlq' AND concurrency_key = '') AS dlq_delta, + sum(throttled) AS throttled_count, + max(queued) AS max_queued, + max(running) AS max_running, + max(queue_limit) AS max_limit, + max(env_queued) AS max_env_queued, + max(env_running) AS max_env_running, + max(env_limit) AS max_env_limit, + max(ck_backlogged) AS max_ck_backlogged, + max(ck_max_wait_ms) AS max_ck_wait_ms, + sumIf(wait_ms, op = 'started' AND concurrency_key = '') AS wait_ms_sum, + countIf(op = 'started' AND wait_ms > 0 AND concurrency_key = '') AS wait_ms_count, + quantilesStateIf(0.5, 0.9, 0.95, 0.99)(wait_ms, op = 'started' AND wait_ms > 0 AND concurrency_key = '') AS wait_quantiles +FROM trigger_dev.queue_metrics_raw_v1 +GROUP BY organization_id, project_id, environment_id, queue_name, bucket_start; + +CREATE MATERIALIZED VIEW IF NOT EXISTS trigger_dev.queue_metrics_ck_mv_v1 +TO trigger_dev.queue_metrics_ck_v1 AS +SELECT + organization_id, project_id, environment_id, queue_name, concurrency_key, + toStartOfInterval(event_time, INTERVAL 10 SECOND) AS bucket_start, + deltaSumTimestampStateIf(cumulative, order_key, op = 'enqueue') AS enqueue_delta, + deltaSumTimestampStateIf(cumulative, order_key, op = 'started') AS started_delta, + deltaSumTimestampStateIf(cumulative, order_key, op = 'ack') AS ack_delta, + maxIf(queued, op = 'gauge') AS max_queued, + maxIf(running, op = 'gauge') AS max_running, + sumIf(wait_ms, op = 'started') AS wait_ms_sum, + countIf(op = 'started' AND wait_ms > 0) AS wait_ms_count +FROM trigger_dev.queue_metrics_raw_v1 +WHERE concurrency_key != '' +GROUP BY organization_id, project_id, environment_id, queue_name, concurrency_key, bucket_start; diff --git a/internal-packages/clickhouse/src/queueMetrics.ts b/internal-packages/clickhouse/src/queueMetrics.ts index 39576b4a0a3..aa3cf5296d2 100644 --- a/internal-packages/clickhouse/src/queueMetrics.ts +++ b/internal-packages/clickhouse/src/queueMetrics.ts @@ -21,6 +21,8 @@ export const QueueMetricsRawV1Input = z.object({ throttled: z.number().optional(), ck_backlogged: z.number().optional(), ck_max_wait_ms: z.number().optional(), + total_running: z.number().optional(), + total_limit: z.number().optional(), wait_ms: z.number().optional(), cumulative: z.number().optional(), }); diff --git a/internal-packages/dashboard-agent/src/__snapshots__/prompt-prefix.test.ts.snap b/internal-packages/dashboard-agent/src/__snapshots__/prompt-prefix.test.ts.snap index d9dac4034a4..da7e4728ff0 100644 --- a/internal-packages/dashboard-agent/src/__snapshots__/prompt-prefix.test.ts.snap +++ b/internal-packages/dashboard-agent/src/__snapshots__/prompt-prefix.test.ts.snap @@ -4,34 +4,34 @@ exports[`the prefix stays inside its budget > matches the committed measurement { "assistant": { "prompt": { - "chars": 33817, - "estimatedTokens": 8454, + "chars": 33913, + "estimatedTokens": 8478, }, "tools": { - "chars": 51289, + "chars": 51313, "count": 25, - "estimatedTokens": 12822, + "estimatedTokens": 12828, }, "total": { - "chars": 85107, - "estimatedTokens": 21277, - "fingerprint": "9d6fcb19", + "chars": 85227, + "estimatedTokens": 21307, + "fingerprint": "d9b6b8b1", }, }, "code": { "prompt": { - "chars": 36572, - "estimatedTokens": 9143, + "chars": 36668, + "estimatedTokens": 9167, }, "tools": { - "chars": 55744, + "chars": 55768, "count": 29, - "estimatedTokens": 13936, + "estimatedTokens": 13942, }, "total": { - "chars": 92317, - "estimatedTokens": 23079, - "fingerprint": "b71774c0", + "chars": 92437, + "estimatedTokens": 23109, + "fingerprint": "552769c8", }, }, } diff --git a/internal-packages/dashboard-agent/src/prompt-prefix.test.ts b/internal-packages/dashboard-agent/src/prompt-prefix.test.ts index 5117d7f064b..9898f5f8431 100644 --- a/internal-packages/dashboard-agent/src/prompt-prefix.test.ts +++ b/internal-packages/dashboard-agent/src/prompt-prefix.test.ts @@ -111,7 +111,7 @@ describe("the head-start and agent prefixes are the same prefix", () => { // learns they exist and how to aggregate them. it("carries the queue chart recipe", () => { expect(DASHBOARD_AGENT_SYSTEM_PROMPT).toContain("Queue charts"); - for (const table of ["queue_metrics", "queue_metrics_by_key"]) { + for (const table of ["concurrency_metrics", "concurrency_metrics_by_key"]) { expect(DASHBOARD_AGENT_SYSTEM_PROMPT).toContain(table); expect(dashboardAgentToolSchemas.get_query_schema.description).toContain(table); } @@ -130,8 +130,8 @@ describe("the head-start and agent prefixes are the same prefix", () => { * drift. The snapshot below is the itemised diff a reviewer reads. */ const PREFIX_BUDGET = { - assistant: { chars: 85_200, estimatedTokens: 21_300, tools: 25, promptChars: 33_900 }, - code: { chars: 92_400, estimatedTokens: 23_100, tools: 29, promptChars: 36_700 }, + assistant: { chars: 85_300, estimatedTokens: 21_350, tools: 25, promptChars: 34_000 }, + code: { chars: 92_500, estimatedTokens: 23_150, tools: 29, promptChars: 36_700 }, } as const; // Measured with watches on: the biggest prefix a turn can hand the provider. diff --git a/internal-packages/dashboard-agent/src/tool-schemas.ts b/internal-packages/dashboard-agent/src/tool-schemas.ts index e94935f59e5..0d7ba0c8709 100644 --- a/internal-packages/dashboard-agent/src/tool-schemas.ts +++ b/internal-packages/dashboard-agent/src/tool-schemas.ts @@ -155,7 +155,7 @@ export const getErrorSchema = tool({ export const getQuerySchemaSchema = tool({ description: - "Discover the analytics tables and columns you can query with TRQL. Call with no table to list the available tables (runs, metrics, llm_metrics, llm_models) and what each holds; call with a table name to get that table's columns, types, descriptions, and time column. Use this before writing a run_query. The queue tables (queue_metrics, queue_metrics_by_key) are queryable but may be missing from the listing — their columns are in the Queue charts guideline, so don't take their absence here as 'no queue data'.", + "Discover the analytics tables and columns you can query with TRQL. Call with no table to list the available tables (runs, metrics, llm_metrics, llm_models) and what each holds; call with a table name to get that table's columns, types, descriptions, and time column. Use this before writing a run_query. The queue tables (concurrency_metrics, concurrency_metrics_by_key) are queryable but may be missing from the listing — their columns are in the Queue charts guideline, so don't take their absence here as 'no queue data'.", inputSchema: z.object({ ...targetFields, table: z @@ -175,7 +175,7 @@ export const runQuerySchema = tool({ query: z .string() .describe( - "The TRQL query. A read-only SELECT over runs / metrics / llm_metrics / llm_models / queue_metrics / queue_metrics_by_key." + "The TRQL query. A read-only SELECT over runs / metrics / llm_metrics / llm_models / concurrency_metrics / concurrency_metrics_by_key." ), period: z .string() @@ -529,11 +529,11 @@ export const KEY_LIST_PLACEHOLDER = ""; const waitQuantile = (index: number, alias: string) => `round(quantilesMerge(0.5, 0.9, 0.95, 0.99)(wait_quantiles)[${index}]) AS ${alias}`; -const FROM_QUEUE = `FROM queue_metrics WHERE queue = '${QUEUE_TEMPLATE_PLACEHOLDER}' GROUP BY t ORDER BY t`; +const FROM_QUEUE = `FROM concurrency_metrics WHERE queue = '${QUEUE_TEMPLATE_PLACEHOLDER}' GROUP BY t ORDER BY t`; const byKey = (expr: string, alias: string) => ({ - rankQuery: `SELECT concurrency_key, ${expr} AS peak FROM queue_metrics_by_key WHERE queue = '${QUEUE_TEMPLATE_PLACEHOLDER}' GROUP BY concurrency_key ORDER BY peak DESC LIMIT 8`, - query: `SELECT timeBucket() AS t, concurrency_key, ${expr} AS ${alias} FROM queue_metrics_by_key WHERE queue = '${QUEUE_TEMPLATE_PLACEHOLDER}' AND concurrency_key IN (${KEY_LIST_PLACEHOLDER}) GROUP BY t, concurrency_key ORDER BY t`, + rankQuery: `SELECT concurrency_key, ${expr} AS peak FROM concurrency_metrics_by_key WHERE queue = '${QUEUE_TEMPLATE_PLACEHOLDER}' GROUP BY concurrency_key ORDER BY peak DESC LIMIT 8`, + query: `SELECT timeBucket() AS t, concurrency_key, ${expr} AS ${alias} FROM concurrency_metrics_by_key WHERE queue = '${QUEUE_TEMPLATE_PLACEHOLDER}' AND concurrency_key IN (${KEY_LIST_PLACEHOLDER}) GROUP BY t, concurrency_key ORDER BY t`, }); export type QueueChartTemplate = { @@ -609,7 +609,7 @@ You have read-only tools that act as the user against their own account: - get_run_trace: a run's execution timeline (spans, durations, errors) — evidence for an investigation into why it failed, retried, or was slow, not itself the answer. - list_errors: distinct errors in the current environment grouped by fingerprint, with occurrence counts and status (unresolved/resolved/ignored). - get_error: full detail for one error group by its error id, including affected versions and who resolved or ignored it. -- get_query_schema: discover the analytics tables and columns you can query with TRQL (runs, metrics, llm_metrics, llm_models, queue_metrics, queue_metrics_by_key). +- get_query_schema: discover the analytics tables and columns you can query with TRQL (runs, metrics, llm_metrics, llm_models, concurrency_metrics, concurrency_metrics_by_key). - run_query: run a read-only TRQL query (SQL-style over ClickHouse) against the current environment's analytics data. - ask_support: ask the Trigger.dev support assistant about how Trigger.dev works (docs, concepts, features, configuration, how-tos). - render_view: render a structured view in the panel from the block catalog. The catalog has the "diagnosis" block (a failure card for a single run), the "chart" block (a line/bar chart of run_query results), the "actions" block (a row of 1-3 buttons offering next steps — a watch intent opens the watch card pre-filled, an ask intent sends the labelled question as the user's next message), and the "investigation" block (a live card for a hypothesis-driven investigation). @@ -708,11 +708,11 @@ Answering with data and charts: - Those buttons are not an offer to do the work: they sit next to a finished answer, and they never license "want me to drill into the top offender?" — asking to look is still banned. Queue charts (the queue page's own charts, as chart blocks): -- queue_metrics is per-queue; queue_metrics_by_key breaks the same queue down per concurrency_key. Time column bucket_start on both, both already scoped to this environment. Pin the queue with WHERE queue = '' (double any quote in the name) and bucket with timeBucket() AS t, which sizes itself from the period. +- concurrency_metrics is per-queue; concurrency_metrics_by_key breaks the same queue down per concurrency_key. Time column bucket_start on both, both already scoped to this environment. Pin the queue with WHERE queue = '' (double any quote in the name) and bucket with timeBucket() AS t, which sizes itself from the period. - Aggregate by column kind or the numbers are wrong: gauges (max_running, max_queued, max_limit, max_ck_backlogged, max_ck_wait_ms) with max(); throttled_count and wait_ms_count with sum(); the counter deltas (enqueue_delta, started_delta, ack_delta) ONLY with deltaSumTimestampMerge(), never sum(); wait latency ONLY with quantilesMerge(0.5, 0.9, 0.95, 0.99)(wait_quantiles)[i], where i is 1 p50, 2 p90, 3 p95, 4 p99. Never FINAL. - The nine charts, all chartType line with xAxisColumn t. Copy the query verbatim and substitute the queue name: ${QUEUE_CHART_RECIPE} -- The last two are two calls: run_query the ranking first, then chart the keys it returned as a quoted IN list. Charting queue_metrics_by_key without that IN list silently loses the newest buckets to the row cap. +- The last two are two calls: run_query the ranking first, then chart the keys it returned as a quoted IN list. Charting concurrency_metrics_by_key without that IN list silently loses the newest buckets to the row cap. - Charts 6-9 only mean something for a queue that shards on concurrencyKey: when max_ck_backlogged is flat zero the queue has no keys, so say that rather than rendering an empty per-key chart.`; // Appended to the system prompt only for turns where watches are enabled: without the diff --git a/internal-packages/database/prisma/migrations/20260827120000_add_task_queue_total_concurrency_limit/migration.sql b/internal-packages/database/prisma/migrations/20260827120000_add_task_queue_total_concurrency_limit/migration.sql new file mode 100644 index 00000000000..dc0ba9a2cac --- /dev/null +++ b/internal-packages/database/prisma/migrations/20260827120000_add_task_queue_total_concurrency_limit/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "TaskQueue" ADD COLUMN IF NOT EXISTS "totalConcurrencyLimit" INTEGER; 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..6ad2b3475b4 --- /dev/null +++ b/internal-packages/database/prisma/migrations/20260829120000_add_queue_gates/migration.sql @@ -0,0 +1,5 @@ +-- AlterTable +ALTER TABLE "BackgroundWorkerTask" ADD COLUMN IF NOT EXISTS "gates" JSONB; + +-- AlterTable +ALTER TABLE "TaskRun" ADD COLUMN IF NOT EXISTS "gates" JSONB; diff --git a/internal-packages/database/prisma/migrations/20260829150000_add_concurrency_overrides/migration.sql b/internal-packages/database/prisma/migrations/20260829150000_add_concurrency_overrides/migration.sql new file mode 100644 index 00000000000..2c7047b0926 --- /dev/null +++ b/internal-packages/database/prisma/migrations/20260829150000_add_concurrency_overrides/migration.sql @@ -0,0 +1,4 @@ +-- AlterTable +ALTER TABLE "TaskQueue" ADD COLUMN IF NOT EXISTS "totalConcurrencyLimitOverriddenAt" TIMESTAMP(3); +ALTER TABLE "TaskQueue" ADD COLUMN IF NOT EXISTS "totalConcurrencyLimitOverriddenBy" TEXT; +ALTER TABLE "TaskQueue" ADD COLUMN IF NOT EXISTS "totalConcurrencyLimitBase" INTEGER; diff --git a/internal-packages/database/prisma/migrations/20260906140000_add_task_queue_concurrency_version_and_role/migration.sql b/internal-packages/database/prisma/migrations/20260906140000_add_task_queue_concurrency_version_and_role/migration.sql new file mode 100644 index 00000000000..eea9fad34b6 --- /dev/null +++ b/internal-packages/database/prisma/migrations/20260906140000_add_task_queue_concurrency_version_and_role/migration.sql @@ -0,0 +1,15 @@ +-- CreateEnum +DO $$ BEGIN + CREATE TYPE "TaskQueueConcurrencyVersion" AS ENUM ('V1', 'V2'); +EXCEPTION WHEN duplicate_object THEN NULL; +END $$; + +-- CreateEnum +DO $$ BEGIN + CREATE TYPE "TaskQueueRole" AS ENUM ('QUEUE', 'LIMIT'); +EXCEPTION WHEN duplicate_object THEN NULL; +END $$; + +-- AlterTable +ALTER TABLE "TaskQueue" ADD COLUMN IF NOT EXISTS "concurrencyVersion" "TaskQueueConcurrencyVersion" NOT NULL DEFAULT 'V1', +ADD COLUMN IF NOT EXISTS "role" "TaskQueueRole" NOT NULL DEFAULT 'QUEUE'; diff --git a/internal-packages/database/prisma/schema.prisma b/internal-packages/database/prisma/schema.prisma index c6f846fa4bd..a93c72c5547 100644 --- a/internal-packages/database/prisma/schema.prisma +++ b/internal-packages/database/prisma/schema.prisma @@ -751,6 +751,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) @@ -1117,6 +1120,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? @@ -1966,6 +1973,13 @@ model TaskQueue { version TaskQueueVersion @default(V1) orderableName String? + /// How the queue's concurrency was declared: V1 = legacy concurrencyLimit semantics + /// (per key when keyed, whole queue otherwise), V2 = the explicit perKey/total shape. + concurrencyVersion TaskQueueConcurrencyVersion @default(V1) + /// QUEUE rows are real queues; LIMIT rows back named concurrency limits and are + /// excluded from every queue-facing surface. + role TaskQueueRole @default(QUEUE) + project Project @relation(fields: [projectId], references: [id], onDelete: Cascade, onUpdate: Cascade) projectId String @@ -1984,6 +1998,15 @@ model TaskQueue { /// percentage (the source of truth). The absolute concurrencyLimit is materialized from it. /// Decimal(5,2) allows fractional percentages like 12.50% (0.01–100.00). concurrencyLimitOverridePercent Decimal? @db.Decimal(5, 2) + /// Caps total concurrent runs across ALL concurrencyKey values of this queue + /// (concurrencyLimit applies per key value). Null = no total cap. + totalConcurrencyLimit Int? + /// When the total concurrency limit was overridden + totalConcurrencyLimitOverriddenAt DateTime? + /// Who overrode the total concurrency limit (null when overridden via the API) + totalConcurrencyLimitOverriddenBy String? + /// If totalConcurrencyLimit is overridden, the declared value it reverts to on reset + totalConcurrencyLimitBase Int? rateLimit Json? paused Boolean @default(false) @@ -1995,9 +2018,11 @@ model TaskQueue { tasks BackgroundWorkerTask[] workers BackgroundWorker[] + @@unique([runtimeEnvironmentId, name]) } + enum TaskQueueType { VIRTUAL NAMED @@ -2008,6 +2033,16 @@ enum TaskQueueVersion { V2 } +enum TaskQueueConcurrencyVersion { + V1 + V2 +} + +enum TaskQueueRole { + QUEUE + LIMIT +} + model BatchTaskRun { id String @id @default(cuid()) friendlyId String @unique diff --git a/internal-packages/metrics-pipeline/src/consumer.test.ts b/internal-packages/metrics-pipeline/src/consumer.test.ts index 672fa426999..f9f59335249 100644 --- a/internal-packages/metrics-pipeline/src/consumer.test.ts +++ b/internal-packages/metrics-pipeline/src/consumer.test.ts @@ -43,6 +43,7 @@ redisTest( }); await consumer.start(); + await emitter.waitUntilReady(); emitter.emit("queueA", { op: "enqueue", q: "queueA" }); emitter.emit("queueB", { op: "started", q: "queueB", wait: 42 }); @@ -156,6 +157,7 @@ redisTest( }); await consumer.start(); + await emitter.waitUntilReady(); emitter.emit(a, { op: "enqueue", q: a }); emitter.emit(b, { op: "enqueue", q: b }); await waitFor(() => inserted.flatMap((i) => i.rows).length >= 2); diff --git a/internal-packages/metrics-pipeline/src/lua.ts b/internal-packages/metrics-pipeline/src/lua.ts index 64f3b896c0d..701f608308a 100644 --- a/internal-packages/metrics-pipeline/src/lua.ts +++ b/internal-packages/metrics-pipeline/src/lua.ts @@ -17,6 +17,10 @@ export type GaugeComputeLuaParams = { // CK-health extras (both or neither): appended as an optional gauge tail, gauge[8]/gauge[9]. ckBacklogged?: string; ckMaxWaitMs?: string; + // Total-concurrency extras (both or neither, and only with the CK extras): appended as + // gauge[10]/gauge[11]. totalLimit is the RAW stored limit (0 = none); readers clamp. + totalRunning?: string; + totalLimit?: string; }; // Computes an op=gauge snapshot into the enclosing script's `__qm_g` local (a flat @@ -26,11 +30,21 @@ export type GaugeComputeLuaParams = { export function createMetricsGaugeComputeLua(params: GaugeComputeLuaParams): string { const throttled = params.throttledExpr ?? "__cc >= __lim and __ql > 0"; const hasCk = params.ckBacklogged != null && params.ckMaxWaitMs != null; - const gauge = hasCk + const hasTotal = params.totalRunning != null && params.totalLimit != null; + if (hasTotal && !hasCk) { + throw new Error("gauge totalRunning/totalLimit extras require the CK extras"); + } + const gauge = hasTotal ? ` local __ckq = tonumber(${params.ckBacklogged}) or 0 local __ckw = tonumber(${params.ckMaxWaitMs}) or 0 + local __tcc = tonumber(${params.totalRunning}) or 0 + local __tlim = tonumber(${params.totalLimit}) or 0 + __qm_g = {__ql, __cc, __lim, __eql, __ec, __elim, __thr, __ckq, __ckw, __tcc, __tlim}` + : hasCk + ? ` local __ckq = tonumber(${params.ckBacklogged}) or 0 + local __ckw = tonumber(${params.ckMaxWaitMs}) or 0 __qm_g = {__ql, __cc, __lim, __eql, __ec, __elim, __thr, __ckq, __ckw}` - : ` __qm_g = {__ql, __cc, __lim, __eql, __ec, __elim, __thr}`; + : ` __qm_g = {__ql, __cc, __lim, __eql, __ec, __elim, __thr}`; return ` if ${params.enabledArg} then 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..c043ab42644 --- /dev/null +++ b/internal-packages/run-engine/src/engine/gateParsing.test.ts @@ -0,0 +1,61 @@ +import { describe, expect, it } from "vitest"; +import { parseGates } from "./gateParsing.js"; + +describe("parseGates", () => { + it("keeps well-shaped gates and caps at four", () => { + expect( + parseGates([ + { queue: "a" }, + { queue: "b", concurrencyKey: "shared" }, + { queue: "c" }, + { queue: "d" }, + { queue: "e" }, + ]) + ).toEqual([ + { queue: "a", concurrencyKey: undefined }, + { queue: "b", concurrencyKey: "shared" }, + { queue: "c", concurrencyKey: undefined }, + { queue: "d", 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..53dccd5bd77 --- /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 four gates apply: three requested gates plus the task's anonymous + * inline-limit gate. + */ +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, 4); +} diff --git a/internal-packages/run-engine/src/engine/index.ts b/internal-packages/run-engine/src/engine/index.ts index fd7cfa43bfb..b28ba9985e4 100644 --- a/internal-packages/run-engine/src/engine/index.ts +++ b/internal-packages/run-engine/src/engine/index.ts @@ -225,6 +225,8 @@ export class RunEngine { queueSelectionStrategy: new FairQueueSelectionStrategy(queueSelectionStrategyOptions), defaultEnvConcurrency: options.queue?.defaultEnvConcurrency ?? 10, defaultEnvConcurrencyBurstFactor: options.queue?.defaultEnvConcurrencyBurstFactor, + totalConcurrencyEnabled: options.queue?.totalConcurrencyEnabled, + gatesEnabled: options.queue?.gatesEnabled, logger: new Logger("RunQueue", options.queue?.logLevel ?? "info"), redis: { ...options.queue.redis, keyPrefix: `${options.queue.redis.keyPrefix}runqueue:` }, retryOptions: options.queue?.retryOptions, @@ -882,6 +884,7 @@ export class RunEngine { sdkVersion, cliVersion, concurrencyKey, + gates, workerQueue, region, enableFastPath, @@ -1063,6 +1066,7 @@ export class RunEngine { sdkVersion, cliVersion, concurrencyKey, + gates, queue, lockedQueueId, workerQueue, @@ -1829,6 +1833,27 @@ export class RunEngine { return this.runQueue.currentConcurrencyOfQueues(environment, queues); } + async totalConcurrencyOfQueues( + environment: MinimalAuthenticatedEnvironment, + queues: string[] + ): Promise> { + return this.runQueue.totalConcurrencyOfQueues(environment, queues); + } + + async totalConcurrencyLimitsOfQueues( + environment: MinimalAuthenticatedEnvironment, + queues: string[] + ): Promise> { + return this.runQueue.totalConcurrencyLimitsOfQueues(environment, queues); + } + + async gateQueuedCountOfQueues( + environment: MinimalAuthenticatedEnvironment, + queues: string[] + ): Promise> { + return this.runQueue.gateQueuedCountOfQueues(environment, queues); + } + async concurrencyKeyBreakdown( environment: MinimalAuthenticatedEnvironment, queue: string, @@ -3107,6 +3132,7 @@ export class RunEngine { { select: { queue: true, + concurrencyKey: true, }, }, this.prisma @@ -3128,6 +3154,7 @@ export class RunEngine { runId, orgId: latestSnapshot.organizationId, queue: taskRun.queue, + concurrencyKey: taskRun.concurrencyKey ?? undefined, env: { id: latestSnapshot.environmentId, type: latestSnapshot.environmentType, diff --git a/internal-packages/run-engine/src/engine/systems/enqueueSystem.ts b/internal-packages/run-engine/src/engine/systems/enqueueSystem.ts index 289e3b9db50..ff4d4a35d0f 100644 --- a/internal-packages/run-engine/src/engine/systems/enqueueSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/enqueueSystem.ts @@ -14,6 +14,7 @@ import { snapshotRouteFromWire, toWireRoute } 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"; @@ -22,6 +23,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; @@ -212,6 +220,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 989da23b690..8bf94d73e36 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 @@ -91,6 +92,10 @@ export type RunEngineOptions = { defaultEnvConcurrency?: number; defaultEnvConcurrencyBurstFactor?: number; logLevel?: LogLevel; + /** Enforce per-queue total concurrency limits across concurrency-key variants. See RunQueueOptions.totalConcurrencyEnabled. */ + totalConcurrencyEnabled?: boolean; + /** Enforce the gates carried in message payloads. See RunQueueOptions.gatesEnabled. */ + gatesEnabled?: boolean; /** Optional queue-metrics emitter; enables gauge + counter emission from the RunQueue. */ queueMetrics?: RunQueueMetricsEmitter; queueSelectionStrategyOptions?: Pick< @@ -332,6 +337,9 @@ export type TriggerParams = { sdkVersion?: string; cliVersion?: string; concurrencyKey?: string; + /** Other queues this run must also hold a concurrency slot in while executing. At + * most four: three requested gates plus the task's anonymous inline-limit gate. */ + 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 d29dc6a8029..f1abdc43ecb 100644 --- a/internal-packages/run-engine/src/index.ts +++ b/internal-packages/run-engine/src/index.ts @@ -115,3 +115,4 @@ export { } from "./engine/waitpointCoordinator/keys.js"; export { createCompletedWaitpointResolver } from "./engine/systems/completedWaitpointResolver.js"; export { enhanceExecutionSnapshotWithWaitpoints } from "./engine/systems/executionSnapshotSystem.js"; +export { parseGates, type ParsedGate } from "./engine/gateParsing.js"; diff --git a/internal-packages/run-engine/src/run-queue/index.ts b/internal-packages/run-engine/src/run-queue/index.ts index 546b61dc229..52511547ab9 100644 --- a/internal-packages/run-engine/src/run-queue/index.ts +++ b/internal-packages/run-engine/src/run-queue/index.ts @@ -58,6 +58,151 @@ const SemanticAttributes = { ORG_ID: "runqueue.orgId", }; +/** + * Gate helpers, spliced into scripts that admit or release runs. A gate is another + * queue named in the message payload that the run must also hold a slot in while it + * executes (payload shape: gates = [{queue, concurrencyKey?}]). Keys are built from + * the payload's own org/proj/env, so gates always live in the message's hash slot. + * + * __gatesHaveCapacity / __gatesAcquire run on admit paths behind the gatesEnabled + * flag. __gatesRelease runs on every release path UNCONDITIONALLY and is payload- + * driven (a cheap substring probe before decoding), so slots acquired while the flag + * was on always drain, and gateless messages pay near zero. Every holder joins the + * gate's groupConcurrency set — keyed or keyless — so a totalConcurrency limit caps + * across everything, matching the SDK's documented "total" semantics. __gateReconcile + * is the same bounded self-heal as the total-cap gate. Group and gate sets are strict + * mirrors of the member's home-queue currentConcurrency set (admits populate both + * in one script), so a member whose message key is gone, or who is absent from its + * home currentConcurrency set, holds no legitimate slot and is pruned. The home-set + * rule is what heals leaks from release paths without the mirror (older builds + * during a rolling upgrade), including runs parked in the DLQ or suspended on + * checkpoints, which older rules based on the queue zset could never prune. + */ +/** TTL for gate queued counters: refreshed on every delta and on every read, so an + * active or observed gate never re-anchors; the Lua helper inlines the same value. */ +const GATE_QUEUED_COUNTER_TTL_SECONDS = 86400; + +const QUEUE_GATES_LUA_HELPERS = ` +local function __gateKeys(gatesKeyPrefix, msg, gate) + local base = gatesKeyPrefix .. '{org:' .. msg.orgId .. '}:proj:' .. msg.projectId .. ':env:' .. msg.environmentId .. ':queue:' .. gate.queue + local gateKey = gate.concurrencyKey + if (not gateKey or gateKey == '') and msg.concurrencyKey and msg.concurrencyKey ~= '' then + gateKey = msg.concurrencyKey + end + local variant = base + if gateKey and gateKey ~= '' then + variant = base .. ':ck:' .. gateKey + end + return base, variant, gateKey +end + +local function __gateReconcile(setKey, msgKeyPrefix, reconcileKeyPrefix) + if not msgKeyPrefix then return end + if redis.call('SET', setKey .. ':reconcileLock', '1', 'NX', 'EX', '10') then + local cursorKey = setKey .. ':reconcileCursor' + local cursor = redis.call('GET', cursorKey) or '0' + local scanResult = redis.call('SSCAN', setKey, cursor, 'COUNT', '100') + redis.call('SET', cursorKey, scanResult[1], 'EX', '3600') + for _, memberId in ipairs(scanResult[2]) do + local rawMemberPayload = redis.call('GET', msgKeyPrefix .. memberId) + if not rawMemberPayload then + redis.call('SREM', setKey, memberId) + elseif reconcileKeyPrefix then + local okMember, member = pcall(cjson.decode, rawMemberPayload) + if okMember and type(member) == 'table' and type(member.queue) == 'string' then + local homeConcurrencyKey = reconcileKeyPrefix .. member.queue .. ':currentConcurrency' + if homeConcurrencyKey ~= setKey and redis.call('SISMEMBER', homeConcurrencyKey, memberId) == 0 then + redis.call('SREM', setKey, memberId) + end + end + end + end + end +end + +local function __gatesHaveCapacity(gatesKeyPrefix, msg, messageId, envLimit, msgKeyPrefix) + if not msg.gates then return true end + for _, gate in ipairs(msg.gates) do + local base, variant, gateKey = __gateKeys(gatesKeyPrefix, msg, gate) + local occupancy = tonumber(redis.call('SCARD', variant .. ':currentConcurrency') or '0') + local perKeyLimit = math.min(tonumber(redis.call('GET', base .. ':concurrency') or '1000000'), envLimit) + if occupancy >= perKeyLimit and redis.call('SISMEMBER', variant .. ':currentConcurrency', messageId) == 0 then + __gateReconcile(variant .. ':currentConcurrency', msgKeyPrefix, gatesKeyPrefix) + return false + end + local rawTotal = redis.call('GET', base .. ':totalConcurrency') + if rawTotal then + local totalLimit = math.min(tonumber(rawTotal), envLimit) + local groupKey = base .. ':groupConcurrency' + if tonumber(redis.call('SCARD', groupKey) or '0') >= totalLimit and redis.call('SISMEMBER', groupKey, messageId) == 0 then + __gateReconcile(groupKey, msgKeyPrefix, gatesKeyPrefix) + return false + end + end + end + return true +end + +local function __gatesAcquire(gatesKeyPrefix, msg, messageId) + if not msg.gates then return end + for _, gate in ipairs(msg.gates) do + local base, variant = __gateKeys(gatesKeyPrefix, msg, gate) + redis.call('SADD', variant .. ':currentConcurrency', messageId) + redis.call('SADD', base .. ':groupConcurrency', messageId) + end +end + +local function __gatesRelease(gatesKeyPrefix, rawPayload, messageId) + if not rawPayload or rawPayload == false then return end + if not string.find(rawPayload, '"gates"', 1, true) then return end + local ok, msg = pcall(cjson.decode, rawPayload) + if not ok or type(msg) ~= 'table' or not msg.gates then return end + for _, gate in ipairs(msg.gates) do + local base, variant = __gateKeys(gatesKeyPrefix, msg, gate) + local removed = redis.call('SREM', variant .. ':currentConcurrency', messageId) + if removed == 1 then + redis.call('SREM', base .. ':groupConcurrency', messageId) + end + end +end + +-- Per-gate queued counter: runs that are queued and must clear the gate to execute. +-- Callers gate the delta on the actual queue-zset transition (ZADD added == 1 / +-- ZREM removed == 1) so re-enqueues and already-removed members never double count; +-- gates sharing a base (duplicate entries, key variants) count once per run. The +-- 24h TTL refreshes on every delta, so an ACTIVE gate's count never resets while +-- drift from delta-less paths (a mixed-version rollout, a stale-entry cleanup) +-- clears once the gate has been quiet for a day. The residual gap is a gate idle +-- for 24h with runs still queued (e.g. paused with no new enqueues): its counter +-- expires and under-counts until the backlog fully drains. Payload-driven and +-- flag-independent, like release, so counts stay exact across flag flips. +local function __gateQueuedDelta(gatesKeyPrefix, msg, delta) + if type(msg) ~= 'table' or not msg.gates then return end + local seenBases = {} + for _, gate in ipairs(msg.gates) do + local base = __gateKeys(gatesKeyPrefix, msg, gate) + if not seenBases[base] then + seenBases[base] = true + local counterKey = base .. ':gateQueuedCounter' + if delta > 0 then + redis.call('INCRBY', counterKey, delta) + redis.call('EXPIRE', counterKey, '86400') + elseif tonumber(redis.call('GET', counterKey) or '0') > 0 then + redis.call('DECRBY', counterKey, -delta) + redis.call('EXPIRE', counterKey, '86400') + end + end + end +end + +local function __gateQueuedDeltaRaw(gatesKeyPrefix, rawPayload, delta) + if not rawPayload or rawPayload == false then return end + if not string.find(rawPayload, '"gates"', 1, true) then return end + local ok, msg = pcall(cjson.decode, rawPayload) + if not ok then return end + __gateQueuedDelta(gatesKeyPrefix, msg, delta) +end`; + // Prelude spliced at the top of every gauge-carrying script: declares the gauge slot and // the return wrapper. A splice fills __qm_g; every return goes through __qmret so the reply // is always {original, gauge}. A nil original becomes false, else Lua drops it from the @@ -66,9 +211,26 @@ const QUEUE_METRICS_GAUGE_PRELUDE = ` local __qm_g = false local function __qmret(r) if r == nil then r = false end return {r, __qm_g} end`; +/** Total-concurrency tail (gauge[10]/gauge[11]): live group cardinality + raw stored cap. + * Requires the groupConcurrencyKey local and the __totalLimitRaw memo (one GET shared with + * the total-cap gate); every script that runs a gauge with this tail declares both. The + * group SCARD stays a fresh read: it must be post-admission. */ +const QUEUE_METRICS_TOTAL_GAUGE_EXTRAS = { + totalRunning: "redis.call('SCARD', groupConcurrencyKey)", + totalLimit: "__totalLimitRaw() or '0'", +}; + +/** The gauge layout is positional (totals ride behind the CK slots), so the plain + * scripts zero-fill the CK health fields: a base queue has no CK variants to backlog. */ +const QUEUE_METRICS_PLAIN_CK_ZERO_EXTRAS = { + ckBacklogged: "0", + ckMaxWaitMs: "0", +}; + // Fresh-read gauge for splice points with no reusable locals: enqueue slow-path (before -// return 0) and the base dequeue top. Gated on the last ARGV so it is inert unless the -// caller opts in. CK queues emit per-subqueue depth (queue_name aggregates via the MV). +// return 0) and the base dequeue's sample-at-return wrapper. Gated on the last ARGV so it +// is inert unless the caller opts in. CK queues emit per-subqueue depth (queue_name +// aggregates via the MV). const QUEUE_METRICS_GAUGE_LUA = createMetricsGaugeComputeLua({ enabledArg: "ARGV[#ARGV] == '1'", queued: "redis.call('ZCARD', queueKey)", @@ -77,6 +239,8 @@ const QUEUE_METRICS_GAUGE_LUA = createMetricsGaugeComputeLua({ envQueued: "redis.call('ZCARD', envQueueKey)", envRunning: "redis.call('SCARD', envCurrentConcurrencyKey)", envLimit: "redis.call('GET', envConcurrencyLimitKey) or defaultEnvConcurrencyLimit", + ...QUEUE_METRICS_PLAIN_CK_ZERO_EXTRAS, + ...QUEUE_METRICS_TOTAL_GAUGE_EXTRAS, }); // Enqueue fast-path gauge: the admission check already computed queueCurrent/envCurrent/ @@ -90,6 +254,8 @@ const QUEUE_METRICS_ENQUEUE_FASTPATH_GAUGE_LUA = createMetricsGaugeComputeLua({ envQueued: "redis.call('ZCARD', envQueueKey)", envRunning: "envCurrent", envLimit: "envLimit", + ...QUEUE_METRICS_PLAIN_CK_ZERO_EXTRAS, + ...QUEUE_METRICS_TOTAL_GAUGE_EXTRAS, }); // CK-health extras: distinct backlogged keys + most-starved head-of-line wait (ckIndex scores @@ -112,6 +278,7 @@ const QUEUE_METRICS_CK_ENQUEUE_GAUGE_LUA = createMetricsGaugeComputeLua({ envRunning: "redis.call('SCARD', envCurrentConcurrencyKey)", envLimit: "redis.call('GET', envConcurrencyLimitKey) or defaultEnvConcurrencyLimit", ...QUEUE_METRICS_CK_GAUGE_EXTRAS, + ...QUEUE_METRICS_TOTAL_GAUGE_EXTRAS, }); const QUEUE_METRICS_CK_ENQUEUE_FASTPATH_GAUGE_LUA = createMetricsGaugeComputeLua({ @@ -123,6 +290,7 @@ const QUEUE_METRICS_CK_ENQUEUE_FASTPATH_GAUGE_LUA = createMetricsGaugeComputeLua envRunning: "envCurrent", envLimit: "envLimit", ...QUEUE_METRICS_CK_GAUGE_EXTRAS, + ...QUEUE_METRICS_TOTAL_GAUGE_EXTRAS, }); // CK dequeue: depth/running from the per-base-queue aggregate counters the run-queue already @@ -138,6 +306,7 @@ const QUEUE_METRICS_CK_DEQUEUE_GAUGE_LUA = createMetricsGaugeComputeLua({ envLimit: "redis.call('GET', envConcurrencyLimitKey) or defaultEnvConcurrencyLimit", throttledExpr: "false", ...QUEUE_METRICS_CK_GAUGE_EXTRAS, + ...QUEUE_METRICS_TOTAL_GAUGE_EXTRAS, }); /** Injected queue-metrics stream emitter; all calls are no-ops when metrics are disabled. */ @@ -179,6 +348,39 @@ export type RunQueueOptions = { * CK operation re-anchors from ckIndex. Default: 86400 (24h). */ counterTtlSeconds?: number; + /** + * When true, queues maintain a per-base-queue groupConcurrency SET (total in-flight + * across all key variants AND keyless runs) and enforce the queue's total concurrency + * limit at admit time. V2 concurrency semantics (a limit's `total` bound, including + * total-only declarations) are enforced solely through this flag: with it off, a + * total-only limit caps nothing. Default false: admit paths are byte-identical to before, and + * only the release-side SREM mirror runs (a no-op on an absent set), so the flag can + * be flipped on a fleet that has fully rolled onto this build without draining queues. + * + * Runs already in flight when the flag turns on are not in the set, so a queue can + * transiently exceed its total limit by the number of those runs. The excess is + * one-time and self-corrects as each pre-flag run completes (its release mirror + * no-ops), after which the limit is enforced exactly. + * + * A release path that misses the group mirror (an instance on an older build during + * rollout) leaves the member behind, briefly under-admitting. The dequeue gate + * reconciles: when a queue sits at its total, members whose message key no longer + * exists or who are absent from their home currentConcurrency set are pruned, so + * such leaks clear within seconds instead of blocking the queue, including runs + * that dead-lettered or suspended through a mirror-less path. Enabling only after + * every instance runs this build avoids the noise but is no longer load-bearing + * for correctness. + */ + totalConcurrencyEnabled?: boolean; + /** + * When true, admit paths enforce the gates carried in a message's payload: the run + * must also have capacity in, and holds a slot in, each gate queue until release. + * Default false: gated messages admit as if they had no gates. Release-side gate + * removal is payload-driven and always on, so slots acquired while the flag was on + * drain correctly after it is turned off, with the same message-key self-heal as + * the total cap covering releases from builds without the mirror. + */ + gatesEnabled?: boolean; workerOptions?: { pollIntervalMs?: number; immediatePollIntervalMs?: number; @@ -464,6 +666,116 @@ export class RunQueue { return result ? Number(result) : undefined; } + public async updateQueueTotalConcurrencyLimits( + env: MinimalAuthenticatedEnvironment, + queue: string, + totalConcurrency: number + ) { + return this.redis.set(this.keys.queueTotalConcurrencyLimitKey(env, queue), totalConcurrency); + } + + public async removeQueueTotalConcurrencyLimits( + env: MinimalAuthenticatedEnvironment, + queue: string + ) { + return this.redis.del(this.keys.queueTotalConcurrencyLimitKey(env, queue)); + } + + public async getQueueTotalConcurrencyLimit(env: MinimalAuthenticatedEnvironment, queue: string) { + const result = await this.redis.get(this.keys.queueTotalConcurrencyLimitKey(env, queue)); + + return result ? Number(result) : undefined; + } + + /** + * Total in-flight runs on a queue, keyed and keyless together (the + * groupConcurrency SET cardinality). Admits only populate the set while + * totalConcurrencyEnabled is on. After the flag is turned off the set drains + * to zero through the release-side mirrors, so a nonzero read reflects real + * runs admitted while it was on, never stale state. + */ + public async totalConcurrencyOfQueue(env: MinimalAuthenticatedEnvironment, queue: string) { + return this.redis.scard(this.keys.queueGroupConcurrencyKey(env, queue)); + } + + /** + * Runs that are queued and must clear this gate queue to execute (the per-gate + * queued counter): incremented per gate on enqueue and decremented on admit and + * on every queued-removal path, floored at zero. Reads refresh the counter's + * TTL, so a gate anyone observes (dashboard, API) never expires while idle — + * e.g. a paused limit with a stalled backlog keeps its count over a quiet + * weekend; only gates nobody touches or reads for a day re-anchor. + */ + public async gateQueuedCountOfQueue(env: MinimalAuthenticatedEnvironment, queue: string) { + const counts = await this.gateQueuedCountOfQueues(env, [queue]); + return counts[queue] ?? 0; + } + + /** Batch variant: one pipeline of GETs, each with a TTL refresh. */ + public async gateQueuedCountOfQueues( + env: MinimalAuthenticatedEnvironment, + queues: string[] + ): Promise> { + const pipeline = this.redis.pipeline(); + queues.forEach((queue) => { + const key = this.keys.gateQueuedCounterKey(env, queue); + pipeline.get(key); + pipeline.expire(key, GATE_QUEUED_COUNTER_TTL_SECONDS, "XX"); + }); + + const results = await pipeline.exec(); + + return queues.reduce( + (acc, queue, index) => { + const value = results?.[index * 2]?.[1]; + const parsed = typeof value === "string" ? Number(value) : 0; + acc[queue] = Number.isFinite(parsed) ? Math.max(parsed, 0) : 0; + return acc; + }, + {} as Record + ); + } + + /** Batch variant of totalConcurrencyOfQueue: one pipeline of group SCARDs. */ + public async totalConcurrencyOfQueues( + env: MinimalAuthenticatedEnvironment, + queues: string[] + ): Promise> { + const pipeline = this.redis.pipeline(); + queues.forEach((queue) => { + pipeline.scard(this.keys.queueGroupConcurrencyKey(env, queue)); + }); + + const results = await pipeline.exec(); + + return queues.reduce( + (acc, queue, index) => { + const value = results?.[index]?.[1]; + acc[queue] = typeof value === "number" ? value : 0; + return acc; + }, + {} as Record + ); + } + + /** Batch read of the RAW stored total concurrency limits (undefined = no cap). */ + public async totalConcurrencyLimitsOfQueues( + env: MinimalAuthenticatedEnvironment, + queues: string[] + ): Promise> { + const keys = queues.map((queue) => this.keys.queueTotalConcurrencyLimitKey(env, queue)); + const values = keys.length > 0 ? await this.redis.mget(...keys) : []; + + return queues.reduce( + (acc, queue, index) => { + const value = values[index]; + acc[queue] = value != null ? Number(value) : undefined; + return acc; + }, + {} as Record + ); + } + public async updateEnvConcurrencyLimits(env: MinimalAuthenticatedEnvironment) { await this.#callUpdateEnvironmentConcurrencyLimits({ envConcurrencyLimitKey: this.keys.envConcurrencyLimitKey(env), @@ -1263,6 +1575,8 @@ export class RunQueue { this.keys.envCurrentDequeuedKeyFromQueue(message.queue), this.keys.queueRunningCounterKeyFromQueue(message.queue), this.keys.ckIndexKeyFromQueue(message.queue), + this.keys.queueGroupConcurrencyKeyFromQueue(message.queue), + this.keys.messageKey(message.orgId, messageId), messageId, this.options.redis.keyPrefix ?? "", String(this.counterTtlSeconds) @@ -1274,7 +1588,10 @@ export class RunQueue { this.keys.envCurrentConcurrencyKeyFromQueue(message.queue), this.keys.queueCurrentDequeuedKeyFromQueue(message.queue), this.keys.envCurrentDequeuedKeyFromQueue(message.queue), - messageId + this.keys.messageKey(message.orgId, messageId), + this.keys.queueGroupConcurrencyKeyFromQueue(message.queue), + messageId, + this.options.redis.keyPrefix ?? "" ); }, { @@ -1331,6 +1648,7 @@ export class RunQueue { runId: string; orgId: string; queue: string; + concurrencyKey?: string; env: RunQueueKeyProducerEnvironment; }) { return this.#callClearMessageFromConcurrencySets(params); @@ -2118,6 +2436,10 @@ export class RunQueue { fields.ckq = ckq; fields.ckw = ckw; } + if (gauge.length >= 11) { + fields.tcc = gauge[9]; + fields.tlim = gauge[10]; + } this.options.queueMetrics?.emitGauge(queue, fields); } @@ -2216,6 +2538,11 @@ export class RunQueue { const lengthCounterKey = this.keys.queueLengthCounterKeyFromQueue(message.queue); const baseQueueKey = this.keys.baseQueueKeyFromQueue(message.queue); const ckKeyPrefix = this.options.redis.keyPrefix ?? ""; + const groupConcurrencyKey = this.keys.queueGroupConcurrencyKeyFromQueue(message.queue); + const totalConcurrencyLimitKey = this.keys.queueTotalConcurrencyLimitKeyFromQueue( + message.queue + ); + const totalConcurrencyEnabledArg = this.options.totalConcurrencyEnabled ? "1" : "0"; if (ttlInfo) { result = await this.redis.enqueueMessageWithTtlCkTracked( @@ -2236,6 +2563,8 @@ export class RunQueue { envConcurrencyLimitBurstFactorKey, lengthCounterKey, baseQueueKey, + groupConcurrencyKey, + totalConcurrencyLimitKey, // args queueName, messageId, @@ -2251,6 +2580,8 @@ export class RunQueue { enableFastPathArg, ckKeyPrefix, String(this.counterTtlSeconds), + totalConcurrencyEnabledArg, + this.options.gatesEnabled ? "1" : "0", metricsGaugeArg ); } else { @@ -2271,6 +2602,8 @@ export class RunQueue { envConcurrencyLimitBurstFactorKey, lengthCounterKey, baseQueueKey, + groupConcurrencyKey, + totalConcurrencyLimitKey, // args queueName, messageId, @@ -2284,6 +2617,8 @@ export class RunQueue { enableFastPathArg, ckKeyPrefix, String(this.counterTtlSeconds), + totalConcurrencyEnabledArg, + this.options.gatesEnabled ? "1" : "0", metricsGaugeArg ); } @@ -2304,6 +2639,8 @@ export class RunQueue { queueConcurrencyLimitKey, envConcurrencyLimitKey, envConcurrencyLimitBurstFactorKey, + this.keys.queueGroupConcurrencyKeyFromQueue(message.queue), + this.keys.queueTotalConcurrencyLimitKeyFromQueue(message.queue), // args queueName, messageId, @@ -2316,6 +2653,9 @@ export class RunQueue { defaultEnvConcurrencyBurstFactor, currentTime, enableFastPathArg, + this.options.redis.keyPrefix ?? "", + this.options.gatesEnabled ? "1" : "0", + this.options.totalConcurrencyEnabled ? "1" : "0", metricsGaugeArg ); } else { @@ -2333,6 +2673,8 @@ export class RunQueue { queueConcurrencyLimitKey, envConcurrencyLimitKey, envConcurrencyLimitBurstFactorKey, + this.keys.queueGroupConcurrencyKeyFromQueue(message.queue), + this.keys.queueTotalConcurrencyLimitKeyFromQueue(message.queue), // args queueName, messageId, @@ -2343,6 +2685,9 @@ export class RunQueue { defaultEnvConcurrencyBurstFactor, currentTime, enableFastPathArg, + this.options.redis.keyPrefix ?? "", + this.options.gatesEnabled ? "1" : "0", + this.options.totalConcurrencyEnabled ? "1" : "0", metricsGaugeArg ); } @@ -2416,6 +2761,8 @@ export class RunQueue { envQueueKey, masterQueueKey, ttlQueueKey, + this.keys.queueGroupConcurrencyKeyFromQueue(messageQueue), + this.keys.queueTotalConcurrencyLimitKeyFromQueue(messageQueue), //args messageQueue, String(Date.now()), @@ -2423,6 +2770,8 @@ export class RunQueue { String(this.options.defaultEnvConcurrencyBurstFactor ?? 1), this.options.redis.keyPrefix ?? "", String(maxCount), + this.options.gatesEnabled ? "1" : "0", + this.options.totalConcurrencyEnabled ? "1" : "0", metricsGaugeArg ); @@ -2549,6 +2898,8 @@ export class RunQueue { ttlQueueKey, lengthCounterKey, runningCounterKey, + this.keys.queueGroupConcurrencyKeyFromQueue(ckWildcardQueue), + this.keys.queueTotalConcurrencyLimitKeyFromQueue(ckWildcardQueue), //args ckWildcardQueue, String(Date.now()), @@ -2556,6 +2907,8 @@ export class RunQueue { String(this.options.defaultEnvConcurrencyBurstFactor ?? 1), this.options.redis.keyPrefix ?? "", String(maxCount), + this.options.totalConcurrencyEnabled ? "1" : "0", + this.options.gatesEnabled ? "1" : "0", metricsGaugeArg ); @@ -2805,11 +3158,13 @@ export class RunQueue { ckIndexKey, lengthCounterKey, runningCounterKey, + this.keys.queueGroupConcurrencyKeyFromQueue(messageQueue), messageId, messageQueue, messageKeyValue, removeFromWorkerQueue ? "1" : "0", - ckWildcardName + ckWildcardName, + this.options.redis.keyPrefix ?? "" ); } @@ -2823,10 +3178,12 @@ export class RunQueue { envCurrentDequeuedKey, envQueueKey, workerQueueKey, + this.keys.queueGroupConcurrencyKeyFromQueue(message.queue), messageId, messageQueue, messageKeyValue, - removeFromWorkerQueue ? "1" : "0" + removeFromWorkerQueue ? "1" : "0", + this.options.redis.keyPrefix ?? "" ); } @@ -2834,18 +3191,31 @@ export class RunQueue { runId, orgId, queue, + concurrencyKey, env, }: { runId: string; orgId: string; queue: string; + concurrencyKey?: string; env: RunQueueKeyProducerEnvironment; }) { const messageId = runId; const messageKey = this.keys.messageKey(orgId, messageId); - const queueCurrentConcurrencyKey = this.keys.queueCurrentConcurrencyKey(env, queue); + /** + * Callers pass the bare TaskRun queue name plus its concurrencyKey; the run's + * slots live on the ck variant. Both variants mirror the per-base-queue group + * set (keyed and keyless admits populate it); the tracked clear additionally + * maintains the counters that only keyed queues keep. + */ + const fullQueue = concurrencyKey ? this.keys.queueKey(env, queue, concurrencyKey) : queue; + const queueCurrentConcurrencyKey = this.keys.queueCurrentConcurrencyKey( + env, + queue, + concurrencyKey + ); const envCurrentConcurrencyKey = this.keys.envCurrentConcurrencyKey(env); - const queueCurrentDequeuedKey = this.keys.queueCurrentDequeuedKey(env, queue); + const queueCurrentDequeuedKey = this.keys.queueCurrentDequeuedKey(env, queue, concurrencyKey); const envCurrentDequeuedKey = this.keys.envCurrentDequeuedKey(env); this.logger.debug("Calling clearMessageFromConcurrencySets", { @@ -2860,14 +3230,16 @@ export class RunQueue { service: this.name, }); - if (queue.includes(":ck:")) { + if (fullQueue.includes(":ck:")) { return this.redis.clearMessageFromConcurrencySetsTracked( queueCurrentConcurrencyKey, envCurrentConcurrencyKey, queueCurrentDequeuedKey, envCurrentDequeuedKey, - this.keys.queueRunningCounterKeyFromQueue(queue), - this.keys.ckIndexKeyFromQueue(queue), + this.keys.queueRunningCounterKeyFromQueue(fullQueue), + this.keys.ckIndexKeyFromQueue(fullQueue), + this.keys.queueGroupConcurrencyKeyFromQueue(fullQueue), + messageKey, messageId, this.options.redis.keyPrefix ?? "", String(this.counterTtlSeconds) @@ -2879,7 +3251,10 @@ export class RunQueue { envCurrentConcurrencyKey, queueCurrentDequeuedKey, envCurrentDequeuedKey, - messageId + messageKey, + this.keys.queueGroupConcurrencyKey(env, queue), + messageId, + this.options.redis.keyPrefix ?? "" ); } @@ -2934,6 +3309,7 @@ export class RunQueue { ckIndexKey, lengthCounterKey, runningCounterKey, + this.keys.queueGroupConcurrencyKeyFromQueue(messageQueue), //args messageId, messageQueue, @@ -2954,11 +3330,13 @@ export class RunQueue { queueCurrentDequeuedKey, envCurrentDequeuedKey, envQueueKey, + this.keys.queueGroupConcurrencyKeyFromQueue(message.queue), //args messageId, messageQueue, JSON.stringify(message), - String(messageScore) + String(messageScore), + this.options.redis.keyPrefix ?? "" ); } } @@ -2997,9 +3375,11 @@ export class RunQueue { ckIndexKey, lengthCounterKey, runningCounterKey, + this.keys.queueGroupConcurrencyKeyFromQueue(messageQueue), messageId, messageQueue, - ckWildcardName + ckWildcardName, + this.options.redis.keyPrefix ?? "" ); } else { await this.redis.moveToDeadLetterQueue( @@ -3012,8 +3392,10 @@ export class RunQueue { envCurrentDequeuedKey, envQueueKey, deadLetterQueueKey, + this.keys.queueGroupConcurrencyKeyFromQueue(message.queue), messageId, - messageQueue + messageQueue, + this.options.redis.keyPrefix ?? "" ); } } @@ -3375,7 +3757,7 @@ end // When enableFastPath == '0', the script skips the fast-path check entirely and behaves // identically to the pre-fast-path version (with the addition of returning 0). this.redis.defineCommand("enqueueMessage", { - numberOfKeys: 12, + numberOfKeys: 14, lua: ` local masterQueueKey = KEYS[1] local queueKey = KEYS[2] @@ -3390,6 +3772,16 @@ local workerQueueKey = KEYS[9] local queueConcurrencyLimitKey = KEYS[10] local envConcurrencyLimitKey = KEYS[11] local envConcurrencyLimitBurstFactorKey = KEYS[12] +-- Total-cap keys (KEYS 13-14) +local groupConcurrencyKey = KEYS[13] +local totalConcurrencyLimitKey = KEYS[14] +local __rawTotalLimit = nil +local function __totalLimitRaw() + if __rawTotalLimit == nil then + __rawTotalLimit = redis.call('GET', totalConcurrencyLimitKey) or false + end + return __rawTotalLimit +end local queueName = ARGV[1] local messageId = ARGV[2] @@ -3401,8 +3793,12 @@ local defaultEnvConcurrencyLimit = ARGV[6] local defaultEnvConcurrencyBurstFactor = ARGV[7] local currentTime = ARGV[8] local enableFastPath = ARGV[9] +local keyPrefix = ARGV[10] +local gatesEnabled = ARGV[11] == '1' +local totalConcurrencyEnabled = ARGV[12] == '1' ${QUEUE_METRICS_GAUGE_PRELUDE} +${QUEUE_GATES_LUA_HELPERS} -- Fast path: check if we can skip the queue and go directly to worker queue if enableFastPath == '1' then @@ -3421,12 +3817,44 @@ if enableFastPath == '1' then ) if queueCurrent < queueLimit then - redis.call('SET', messageKey, messageData) - redis.call('SADD', queueCurrentConcurrencyKey, messageId) - redis.call('SADD', envCurrentConcurrencyKey, messageId) - redis.call('RPUSH', workerQueueKey, messageKeyValue) + -- Total-cap gate: a fast-path admit consumes a group slot, so it must + -- respect the env-clamped total limit. At the cap we fall through to the + -- slow path (the message queues; the dequeue gate holds it). + local totalAllowsFastPath = true + if totalConcurrencyEnabled then + local rawTotalLimit = __totalLimitRaw() + if rawTotalLimit then + local totalLimit = math.min(tonumber(rawTotalLimit), envLimit) + if tonumber(redis.call('SCARD', groupConcurrencyKey) or '0') >= totalLimit then + totalAllowsFastPath = false + end + end + end + + local gateMsg = nil + local gatesAllowFastPath = true + if gatesEnabled and string.find(messageData, '"gates"', 1, true) then + local okDecode, decoded = pcall(cjson.decode, messageData) + if okDecode and type(decoded) == 'table' and decoded.gates then + gateMsg = decoded + gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil) + end + end + + if totalAllowsFastPath and gatesAllowFastPath then + redis.call('SET', messageKey, messageData) + redis.call('SADD', queueCurrentConcurrencyKey, messageId) + redis.call('SADD', envCurrentConcurrencyKey, messageId) + if totalConcurrencyEnabled then + redis.call('SADD', groupConcurrencyKey, messageId) + end + if gateMsg then + __gatesAcquire(keyPrefix, gateMsg, messageId) + end + redis.call('RPUSH', workerQueueKey, messageKeyValue) ${QUEUE_METRICS_ENQUEUE_FASTPATH_GAUGE_LUA} - return __qmret(1) + return __qmret(1) + end end end end @@ -3437,7 +3865,10 @@ end redis.call('SET', messageKey, messageData) -- Add the message to the queue -redis.call('ZADD', queueKey, messageScore, messageId) +local added = redis.call('ZADD', queueKey, messageScore, messageId) +if added == 1 then + __gateQueuedDeltaRaw(keyPrefix, messageData, 1) +end -- Add the message to the env queue redis.call('ZADD', envQueueKey, messageScore, messageId) @@ -3451,11 +3882,16 @@ else redis.call('ZADD', masterQueueKey, earliestMessage[2], queueName) end --- Update the concurrency keys -redis.call('SREM', queueCurrentConcurrencyKey, messageId) +-- Update the concurrency keys. The groupConcurrency SREM mirrors the base SREM +-- unconditionally (no flag check) so a disabled flag still drains the group set. +local removedFromCurrentConcurrency = redis.call('SREM', queueCurrentConcurrencyKey, messageId) +if removedFromCurrentConcurrency == 1 then + redis.call('SREM', groupConcurrencyKey, messageId) +end redis.call('SREM', envCurrentConcurrencyKey, messageId) redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) +__gatesRelease(keyPrefix, messageData, messageId) ${QUEUE_METRICS_GAUGE_LUA} return __qmret(0) @@ -3467,7 +3903,7 @@ return __qmret(0) // (scheduled independently before enqueue) handles TTL expiry. This mirrors what // dequeueMessagesFromQueue does: it removes from the TTL set when dequeuing. this.redis.defineCommand("enqueueMessageWithTtl", { - numberOfKeys: 13, + numberOfKeys: 15, lua: ` local masterQueueKey = KEYS[1] local queueKey = KEYS[2] @@ -3483,6 +3919,16 @@ local workerQueueKey = KEYS[10] local queueConcurrencyLimitKey = KEYS[11] local envConcurrencyLimitKey = KEYS[12] local envConcurrencyLimitBurstFactorKey = KEYS[13] +-- Total-cap keys (KEYS 14-15) +local groupConcurrencyKey = KEYS[14] +local totalConcurrencyLimitKey = KEYS[15] +local __rawTotalLimit = nil +local function __totalLimitRaw() + if __rawTotalLimit == nil then + __rawTotalLimit = redis.call('GET', totalConcurrencyLimitKey) or false + end + return __rawTotalLimit +end local queueName = ARGV[1] local messageId = ARGV[2] @@ -3496,8 +3942,12 @@ local defaultEnvConcurrencyLimit = ARGV[8] local defaultEnvConcurrencyBurstFactor = ARGV[9] local currentTime = ARGV[10] local enableFastPath = ARGV[11] +local keyPrefix = ARGV[12] +local gatesEnabled = ARGV[13] == '1' +local totalConcurrencyEnabled = ARGV[14] == '1' ${QUEUE_METRICS_GAUGE_PRELUDE} +${QUEUE_GATES_LUA_HELPERS} -- Fast path: check if we can skip the queue and go directly to worker queue if enableFastPath == '1' then @@ -3516,13 +3966,45 @@ if enableFastPath == '1' then ) if queueCurrent < queueLimit then - redis.call('SET', messageKey, messageData) - redis.call('SADD', queueCurrentConcurrencyKey, messageId) - redis.call('SADD', envCurrentConcurrencyKey, messageId) - redis.call('RPUSH', workerQueueKey, messageKeyValue) + -- Total-cap gate: a fast-path admit consumes a group slot, so it must + -- respect the env-clamped total limit. At the cap we fall through to the + -- slow path (the message queues; the dequeue gate holds it). + local totalAllowsFastPath = true + if totalConcurrencyEnabled then + local rawTotalLimit = __totalLimitRaw() + if rawTotalLimit then + local totalLimit = math.min(tonumber(rawTotalLimit), envLimit) + if tonumber(redis.call('SCARD', groupConcurrencyKey) or '0') >= totalLimit then + totalAllowsFastPath = false + end + end + end + + local gateMsg = nil + local gatesAllowFastPath = true + if gatesEnabled and string.find(messageData, '"gates"', 1, true) then + local okDecode, decoded = pcall(cjson.decode, messageData) + if okDecode and type(decoded) == 'table' and decoded.gates then + gateMsg = decoded + gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil) + end + end + + if totalAllowsFastPath and gatesAllowFastPath then + redis.call('SET', messageKey, messageData) + redis.call('SADD', queueCurrentConcurrencyKey, messageId) + redis.call('SADD', envCurrentConcurrencyKey, messageId) + if totalConcurrencyEnabled then + redis.call('SADD', groupConcurrencyKey, messageId) + end + if gateMsg then + __gatesAcquire(keyPrefix, gateMsg, messageId) + end + redis.call('RPUSH', workerQueueKey, messageKeyValue) ${QUEUE_METRICS_ENQUEUE_FASTPATH_GAUGE_LUA} -- Skip TTL sorted set: the expireRun worker job handles TTL expiry independently - return __qmret(1) + return __qmret(1) + end end end end @@ -3533,7 +4015,10 @@ end redis.call('SET', messageKey, messageData) -- Add the message to the queue -redis.call('ZADD', queueKey, messageScore, messageId) +local added = redis.call('ZADD', queueKey, messageScore, messageId) +if added == 1 then + __gateQueuedDeltaRaw(keyPrefix, messageData, 1) +end -- Add the message to the env queue redis.call('ZADD', envQueueKey, messageScore, messageId) @@ -3550,11 +4035,16 @@ else redis.call('ZADD', masterQueueKey, earliestMessage[2], queueName) end --- Update the concurrency keys -redis.call('SREM', queueCurrentConcurrencyKey, messageId) +-- Update the concurrency keys. The groupConcurrency SREM mirrors the base SREM +-- unconditionally (no flag check) so a disabled flag still drains the group set. +local removedFromCurrentConcurrency = redis.call('SREM', queueCurrentConcurrencyKey, messageId) +if removedFromCurrentConcurrency == 1 then + redis.call('SREM', groupConcurrencyKey, messageId) +end redis.call('SREM', envCurrentConcurrencyKey, messageId) redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) +__gatesRelease(keyPrefix, messageData, messageId) ${QUEUE_METRICS_GAUGE_LUA} return __qmret(0) @@ -3785,7 +4275,7 @@ return __qmret(0) // *Tracked variants of dequeueMessageFromKey and the ack/nack/dlq/release/clear // scripts. this.redis.defineCommand("enqueueMessageCkTracked", { - numberOfKeys: 15, + numberOfKeys: 17, lua: ` local masterQueueKey = KEYS[1] local queueKey = KEYS[2] @@ -3804,6 +4294,16 @@ local envConcurrencyLimitBurstFactorKey = KEYS[13] -- Counter keys (KEYS 14-15) local lengthCounterKey = KEYS[14] local baseQueueKey = KEYS[15] +-- Total-cap keys (KEYS 16-17) +local groupConcurrencyKey = KEYS[16] +local totalConcurrencyLimitKey = KEYS[17] +local __rawTotalLimit = nil +local function __totalLimitRaw() + if __rawTotalLimit == nil then + __rawTotalLimit = redis.call('GET', totalConcurrencyLimitKey) or false + end + return __rawTotalLimit +end local queueName = ARGV[1] local messageId = ARGV[2] @@ -3820,8 +4320,11 @@ local enableFastPath = ARGV[10] local keyPrefix = ARGV[11] -- TTL (seconds) applied to counter lazy-init SETs local counterTtl = ARGV[12] +local totalConcurrencyEnabled = ARGV[13] == '1' +local gatesEnabled = ARGV[14] == '1' ${QUEUE_METRICS_GAUGE_PRELUDE} +${QUEUE_GATES_LUA_HELPERS} -- Fast path: check if we can skip the queue and go directly to worker queue if enableFastPath == '1' then @@ -3840,15 +4343,47 @@ if enableFastPath == '1' then ) if queueCurrent < queueLimit then - redis.call('SET', messageKey, messageData) - redis.call('SADD', queueCurrentConcurrencyKey, messageId) - redis.call('SADD', envCurrentConcurrencyKey, messageId) - redis.call('RPUSH', workerQueueKey, messageKeyValue) + -- Total-cap gate: a fast-path admit consumes a group slot, so it must + -- respect the env-clamped total limit. At the cap we fall through to the + -- slow path (the message queues; the dequeue gate holds it). + local totalAllowsFastPath = true + if totalConcurrencyEnabled then + local rawTotalLimit = __totalLimitRaw() + if rawTotalLimit then + local totalLimit = math.min(tonumber(rawTotalLimit), envLimit) + if tonumber(redis.call('SCARD', groupConcurrencyKey) or '0') >= totalLimit then + totalAllowsFastPath = false + end + end + end + + local gateMsg = nil + local gatesAllowFastPath = true + if gatesEnabled and string.find(messageData, '"gates"', 1, true) then + local okDecode, decoded = pcall(cjson.decode, messageData) + if okDecode and type(decoded) == 'table' and decoded.gates then + gateMsg = decoded + gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil) + end + end + + if totalAllowsFastPath and gatesAllowFastPath then + redis.call('SET', messageKey, messageData) + redis.call('SADD', queueCurrentConcurrencyKey, messageId) + redis.call('SADD', envCurrentConcurrencyKey, messageId) + if totalConcurrencyEnabled then + redis.call('SADD', groupConcurrencyKey, messageId) + end + if gateMsg then + __gatesAcquire(keyPrefix, gateMsg, messageId) + end + redis.call('RPUSH', workerQueueKey, messageKeyValue) ${QUEUE_METRICS_CK_ENQUEUE_FASTPATH_GAUGE_LUA} - -- Fast-path skips the CK variant zset entirely; lengthCounter is unchanged. - -- runningCounter is bumped later by dequeueMessageFromKeyTracked when the - -- worker pulls the message from the worker queue. - return __qmret(1) + -- Fast-path skips the CK variant zset entirely; lengthCounter is unchanged. + -- runningCounter is bumped later by dequeueMessageFromKeyTracked when the + -- worker pulls the message from the worker queue. + return __qmret(1) + end end end end @@ -3878,6 +4413,7 @@ local added = redis.call('ZADD', queueKey, messageScore, messageId) redis.call('ZADD', envQueueKey, messageScore, messageId) if added == 1 then redis.call('INCR', lengthCounterKey) + __gateQueuedDeltaRaw(keyPrefix, messageData, 1) end -- Rebalance CK index @@ -3900,11 +4436,16 @@ if queueName ~= ckWildcardName then redis.call('ZREM', masterQueueKey, queueName) end --- Update the concurrency keys -redis.call('SREM', queueCurrentConcurrencyKey, messageId) +-- Update the concurrency keys. The groupConcurrency SREM mirrors the per-CK SREM +-- unconditionally (no flag check) so a disabled flag still drains the group set. +local removedFromCurrentConcurrency = redis.call('SREM', queueCurrentConcurrencyKey, messageId) +if removedFromCurrentConcurrency == 1 then + redis.call('SREM', groupConcurrencyKey, messageId) +end redis.call('SREM', envCurrentConcurrencyKey, messageId) redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) +__gatesRelease(keyPrefix, messageData, messageId) ${QUEUE_METRICS_CK_ENQUEUE_GAUGE_LUA} return __qmret(0) @@ -3912,7 +4453,7 @@ return __qmret(0) }); this.redis.defineCommand("enqueueMessageWithTtlCkTracked", { - numberOfKeys: 16, + numberOfKeys: 18, lua: ` local masterQueueKey = KEYS[1] local queueKey = KEYS[2] @@ -3932,6 +4473,16 @@ local envConcurrencyLimitBurstFactorKey = KEYS[14] -- Counter keys (KEYS 15-16) local lengthCounterKey = KEYS[15] local baseQueueKey = KEYS[16] +-- Total-cap keys (KEYS 17-18) +local groupConcurrencyKey = KEYS[17] +local totalConcurrencyLimitKey = KEYS[18] +local __rawTotalLimit = nil +local function __totalLimitRaw() + if __rawTotalLimit == nil then + __rawTotalLimit = redis.call('GET', totalConcurrencyLimitKey) or false + end + return __rawTotalLimit +end local queueName = ARGV[1] local messageId = ARGV[2] @@ -3950,8 +4501,11 @@ local enableFastPath = ARGV[12] local keyPrefix = ARGV[13] -- TTL (seconds) applied to counter lazy-init SETs local counterTtl = ARGV[14] +local totalConcurrencyEnabled = ARGV[15] == '1' +local gatesEnabled = ARGV[16] == '1' ${QUEUE_METRICS_GAUGE_PRELUDE} +${QUEUE_GATES_LUA_HELPERS} -- Fast path: check if we can skip the queue and go directly to worker queue if enableFastPath == '1' then @@ -3970,12 +4524,42 @@ if enableFastPath == '1' then ) if queueCurrent < queueLimit then - redis.call('SET', messageKey, messageData) - redis.call('SADD', queueCurrentConcurrencyKey, messageId) - redis.call('SADD', envCurrentConcurrencyKey, messageId) - redis.call('RPUSH', workerQueueKey, messageKeyValue) + -- Total-cap gate: see enqueueMessageCkTracked. + local totalAllowsFastPath = true + if totalConcurrencyEnabled then + local rawTotalLimit = __totalLimitRaw() + if rawTotalLimit then + local totalLimit = math.min(tonumber(rawTotalLimit), envLimit) + if tonumber(redis.call('SCARD', groupConcurrencyKey) or '0') >= totalLimit then + totalAllowsFastPath = false + end + end + end + + local gateMsg = nil + local gatesAllowFastPath = true + if gatesEnabled and string.find(messageData, '"gates"', 1, true) then + local okDecode, decoded = pcall(cjson.decode, messageData) + if okDecode and type(decoded) == 'table' and decoded.gates then + gateMsg = decoded + gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil) + end + end + + if totalAllowsFastPath and gatesAllowFastPath then + redis.call('SET', messageKey, messageData) + redis.call('SADD', queueCurrentConcurrencyKey, messageId) + redis.call('SADD', envCurrentConcurrencyKey, messageId) + if totalConcurrencyEnabled then + redis.call('SADD', groupConcurrencyKey, messageId) + end + if gateMsg then + __gatesAcquire(keyPrefix, gateMsg, messageId) + end + redis.call('RPUSH', workerQueueKey, messageKeyValue) ${QUEUE_METRICS_CK_ENQUEUE_FASTPATH_GAUGE_LUA} - return __qmret(1) + return __qmret(1) + end end end end @@ -4001,6 +4585,7 @@ redis.call('ZADD', envQueueKey, messageScore, messageId) redis.call('ZADD', ttlQueueKey, ttlScore, ttlMember) if added == 1 then redis.call('INCR', lengthCounterKey) + __gateQueuedDeltaRaw(keyPrefix, messageData, 1) end -- Rebalance CK index @@ -4023,11 +4608,16 @@ if queueName ~= ckWildcardName then redis.call('ZREM', masterQueueKey, queueName) end --- Update the concurrency keys -redis.call('SREM', queueCurrentConcurrencyKey, messageId) +-- Update the concurrency keys. The groupConcurrency SREM mirrors the per-CK SREM +-- unconditionally (no flag check) so a disabled flag still drains the group set. +local removedFromCurrentConcurrency = redis.call('SREM', queueCurrentConcurrencyKey, messageId) +if removedFromCurrentConcurrency == 1 then + redis.call('SREM', groupConcurrencyKey, messageId) +end redis.call('SREM', envCurrentConcurrencyKey, messageId) redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) +__gatesRelease(keyPrefix, messageData, messageId) ${QUEUE_METRICS_CK_ENQUEUE_GAUGE_LUA} return __qmret(0) @@ -4174,6 +4764,7 @@ local function decrFloored(key) redis.call('DECR', key) end end +${QUEUE_GATES_LUA_HELPERS} local expiredMembers = redis.call('ZRANGEBYSCORE', ttlQueueKey, '-inf', currentTime, 'LIMIT', 0, batchSize) @@ -4205,12 +4796,13 @@ for i, member in ipairs(expiredMembers) do local messageKey = keyPrefix .. "{org:" .. orgFromQueue .. "}:message:" .. runId - -- Read the message's versioned snapshotRoute BEFORE deleting it, so the TTL worker resolves each + -- Read the message once BEFORE deleting it: the raw payload feeds the gate release and + -- queued-counter delta, and its versioned snapshotRoute lets the TTL worker resolve each -- run's residency from the route the birth stamped (no per-run durable lookup on the worker side). local snapshotRoute = nil - local rawMessage = redis.call('GET', messageKey) - if rawMessage then - local ok, decoded = pcall(cjson.decode, rawMessage) + local rawPayload = redis.call('GET', messageKey) + if rawPayload then + local ok, decoded = pcall(cjson.decode, rawPayload) if ok and type(decoded) == 'table' and decoded.snapshotRoute ~= nil then snapshotRoute = decoded.snapshotRoute end @@ -4220,6 +4812,9 @@ for i, member in ipairs(expiredMembers) do -- ZREM from queue; if successful AND this is a CK variant, DECR lengthCounter. local removedFromZset = redis.call('ZREM', queueKey, runId) + if removedFromZset == 1 then + __gateQueuedDeltaRaw(keyPrefix, rawPayload, -1) + end local envMatch = string.match(rawQueueKey, ":env:([^:]+)") if envMatch then @@ -4229,8 +4824,9 @@ for i, member in ipairs(expiredMembers) do local concurrencyKey = queueKey .. ":currentConcurrency" local dequeuedKey = queueKey .. ":currentDequeued" - redis.call('SREM', concurrencyKey, runId) + local removedFromCurrent = redis.call('SREM', concurrencyKey, runId) local removedFromDequeued = redis.call('SREM', dequeuedKey, runId) + __gatesRelease(keyPrefix, rawPayload, runId) local projMatch = string.match(rawQueueKey, ":proj:([^:]+):env:") local envConcurrencyKey = keyPrefix .. "{org:" .. orgFromQueue .. "}:proj:" .. (projMatch or "") .. ":env:" .. (envMatch or "") .. ":currentConcurrency" @@ -4238,8 +4834,16 @@ for i, member in ipairs(expiredMembers) do redis.call('SREM', envConcurrencyKey, runId) redis.call('SREM', envDequeuedKey, runId) - -- Rebalance CK index AND update counters if this is a CK queue + -- Mirror the currentConcurrency SREM into the base groupConcurrency set for + -- keyed and keyless queues alike, so a defensive removal always drains the + -- total pool too. local ckMatch = string.match(rawQueueKey, "(.-):ck:") + if removedFromCurrent == 1 then + local groupBase = ckMatch or rawQueueKey + redis.call('SREM', keyPrefix .. groupBase .. ":groupConcurrency", runId) + end + + -- Rebalance CK index AND update counters if this is a CK queue if ckMatch then local lengthCounterKey = keyPrefix .. ckMatch .. ":lengthCounter" local runningCounterKey = keyPrefix .. ckMatch .. ":runningCounter" @@ -4278,7 +4882,7 @@ return results }); this.redis.defineCommand("dequeueMessagesFromQueue", { - numberOfKeys: 10, + numberOfKeys: 12, lua: ` local queueKey = KEYS[1] local queueConcurrencyLimitKey = KEYS[2] @@ -4290,6 +4894,16 @@ local messageKeyPrefix = KEYS[7] local envQueueKey = KEYS[8] local masterQueueKey = KEYS[9] local ttlQueueKey = KEYS[10] -- Optional: TTL sorted set key (empty string if not used) +-- Total-cap keys (KEYS 11-12) +local groupConcurrencyKey = KEYS[11] +local totalConcurrencyLimitKey = KEYS[12] +local __rawTotalLimit = nil +local function __totalLimitRaw() + if __rawTotalLimit == nil then + __rawTotalLimit = redis.call('GET', totalConcurrencyLimitKey) or false + end + return __rawTotalLimit +end local queueName = ARGV[1] local currentTime = tonumber(ARGV[2]) @@ -4297,8 +4911,20 @@ local defaultEnvConcurrencyLimit = ARGV[3] local defaultEnvConcurrencyBurstFactor = ARGV[4] local keyPrefix = ARGV[5] local maxCount = tonumber(ARGV[6] or '1') +local gatesEnabled = ARGV[7] == '1' +local totalConcurrencyEnabled = ARGV[8] == '1' ${QUEUE_METRICS_GAUGE_PRELUDE} +${QUEUE_GATES_LUA_HELPERS} +-- Sample-at-return: the gauge is computed once, by the return wrapper, so every +-- exit emits the state as of that exit (post-admission on the success path) and +-- no path pays for a sample that a later one would overwrite. +local function __qmsample() ${QUEUE_METRICS_GAUGE_LUA} +end +do + local __qmret_inner = __qmret + __qmret = function(r) __qmsample() return __qmret_inner(r) end +end -- Check current env concurrency against the limit local envCurrentConcurrency = tonumber(redis.call('SCARD', envCurrentConcurrencyKey) or '0') @@ -4325,6 +4951,22 @@ local envAvailableCapacity = envConcurrencyLimitWithBurstFactor - envCurrentConc local queueAvailableCapacity = totalQueueConcurrencyLimit - queueCurrentConcurrency local actualMaxCount = math.min(maxCount, envAvailableCapacity, queueAvailableCapacity) +-- Total-cap gate: every admit joins the group set, so the batch is bounded by the +-- env-clamped total limit's remaining capacity. At saturation, run the bounded +-- reconcile once (heals leaked members) and re-check before giving up. +if totalConcurrencyEnabled then + local rawTotalLimit = __totalLimitRaw() + if rawTotalLimit then + local totalLimit = math.min(tonumber(rawTotalLimit), envConcurrencyLimit) + local groupCurrentConcurrency = tonumber(redis.call('SCARD', groupConcurrencyKey) or '0') + if groupCurrentConcurrency >= totalLimit then + __gateReconcile(groupConcurrencyKey, messageKeyPrefix, keyPrefix) + groupCurrentConcurrency = tonumber(redis.call('SCARD', groupConcurrencyKey) or '0') + end + actualMaxCount = math.min(actualMaxCount, totalLimit - groupCurrentConcurrency) + end +end + if actualMaxCount <= 0 then return __qmret(nil) end @@ -4359,31 +5001,47 @@ for i = 1, #messages, 2 do -- leave messageKey intact, and (re-)register the TTL entry so the -- TTL consumer can discover and properly expire the run. The entry -- is removed on first dequeue, so it cannot be assumed to exist. - redis.call('ZREM', queueKey, messageId) + local removedExpired = redis.call('ZREM', queueKey, messageId) redis.call('ZREM', envQueueKey, messageId) + if removedExpired == 1 then + __gateQueuedDelta(keyPrefix, messageData, -1) + end if ttlQueueKey and ttlQueueKey ~= '' then local ttlMember = queueName .. '|' .. messageId .. '|' .. (messageData.orgId or '') redis.call('ZADD', ttlQueueKey, ttlExpiresAt, ttlMember) end else - -- Not expired - process normally - redis.call('ZREM', queueKey, messageId) - redis.call('ZREM', envQueueKey, messageId) - redis.call('SADD', queueCurrentConcurrencyKey, messageId) - redis.call('SADD', envCurrentConcurrencyKey, messageId) - - -- Remove from TTL set if provided (run is being executed, not expired) - if ttlQueueKey and ttlQueueKey ~= '' and ttlExpiresAt then - local ttlMember = queueName .. '|' .. messageId .. '|' .. (messageData.orgId or '') - redis.call('ZREM', ttlQueueKey, ttlMember) + local gatesAllow = true + if gatesEnabled then + gatesAllow = __gatesHaveCapacity(keyPrefix, messageData, messageId, envConcurrencyLimit, messageKeyPrefix) end - -- Add to results - table.insert(results, messageId) - table.insert(results, messageScore) - table.insert(results, messagePayload) - - dequeuedCount = dequeuedCount + 1 + if gatesAllow then + local removedFromQueue = redis.call('ZREM', queueKey, messageId) + redis.call('ZREM', envQueueKey, messageId) + if removedFromQueue == 1 then + __gateQueuedDelta(keyPrefix, messageData, -1) + end + redis.call('SADD', queueCurrentConcurrencyKey, messageId) + redis.call('SADD', envCurrentConcurrencyKey, messageId) + if totalConcurrencyEnabled then + redis.call('SADD', groupConcurrencyKey, messageId) + end + if gatesEnabled then + __gatesAcquire(keyPrefix, messageData, messageId) + end + + if ttlQueueKey and ttlQueueKey ~= '' and ttlExpiresAt then + local ttlMember = queueName .. '|' .. messageId .. '|' .. (messageData.orgId or '') + redis.call('ZREM', ttlQueueKey, ttlMember) + end + + table.insert(results, messageId) + table.insert(results, messageScore) + table.insert(results, messagePayload) + + dequeuedCount = dequeuedCount + 1 + end end else -- Stale entry: message key was already deleted (e.g. acknowledged), @@ -4563,7 +5221,7 @@ return results // (normal dequeue, TTL-expired, or stale-orphan path — all of which were // counted at enqueue time). this.redis.defineCommand("dequeueMessagesFromCkQueueTracked", { - numberOfKeys: 11, + numberOfKeys: 13, lua: ` local ckIndexKey = KEYS[1] local queueConcurrencyLimitKey = KEYS[2] @@ -4576,6 +5234,15 @@ local masterQueueKey = KEYS[8] local ttlQueueKey = KEYS[9] local lengthCounterKey = KEYS[10] local runningCounterKey = KEYS[11] +local groupConcurrencyKey = KEYS[12] +local totalConcurrencyLimitKey = KEYS[13] +local __rawTotalLimit = nil +local function __totalLimitRaw() + if __rawTotalLimit == nil then + __rawTotalLimit = redis.call('GET', totalConcurrencyLimitKey) or false + end + return __rawTotalLimit +end local ckWildcardName = ARGV[1] local currentTime = tonumber(ARGV[2]) @@ -4583,8 +5250,20 @@ local defaultEnvConcurrencyLimit = ARGV[3] local defaultEnvConcurrencyBurstFactor = ARGV[4] local keyPrefix = ARGV[5] local maxCount = tonumber(ARGV[6] or '1') +local totalConcurrencyEnabled = ARGV[7] == '1' +local gatesEnabled = ARGV[8] == '1' ${QUEUE_METRICS_GAUGE_PRELUDE} +${QUEUE_GATES_LUA_HELPERS} +-- Sample-at-return: the gauge is computed once, by the return wrapper, so every +-- exit emits the state as of that exit (post-admission on the success path) and +-- no path pays for a sample that a later one would overwrite. +local function __qmsample() ${QUEUE_METRICS_CK_DEQUEUE_GAUGE_LUA} +end +do + local __qmret_inner = __qmret + __qmret = function(r) __qmsample() return __qmret_inner(r) end +end local function decrLengthCounter() if tonumber(redis.call('GET', lengthCounterKey) or '0') > 0 then @@ -4607,6 +5286,32 @@ local queueConcurrencyLimit = math.min(tonumber(redis.call('GET', queueConcurren local envAvailableCapacity = envConcurrencyLimitWithBurstFactor - envCurrentConcurrency local actualMaxCount = math.min(maxCount, envAvailableCapacity) +-- Total-cap gate: track the remaining headroom across ALL ck variants +-- (groupConcurrency SCARD vs the env-clamped total limit). Admission below +-- decrements the headroom per new member, and a message that is ALREADY a +-- group member passes even at zero headroom: it holds its own slot (an +-- unmirrored release from an older build can leave a queued run's membership +-- behind, and blocking on it would deadlock the run against itself). +local totalHeadroom = nil +if totalConcurrencyEnabled then + local rawTotalLimit = __totalLimitRaw() + if rawTotalLimit then + local totalConcurrencyLimit = math.min(tonumber(rawTotalLimit), envConcurrencyLimit) + local groupCurrentConcurrency = tonumber(redis.call('SCARD', groupConcurrencyKey) or '0') + + -- Self-heal before holding the queue at its limit: a member with no message + -- key is dead, and a member absent from its home currentConcurrency set is + -- not in flight (group membership is a strict mirror of it), so neither + -- holds a legitimate slot. Both are pruned by the shared bounded reconcile. + if groupCurrentConcurrency >= totalConcurrencyLimit then + __gateReconcile(groupConcurrencyKey, messageKeyPrefix, keyPrefix) + groupCurrentConcurrency = tonumber(redis.call('SCARD', groupConcurrencyKey) or '0') + end + + totalHeadroom = totalConcurrencyLimit - groupCurrentConcurrency + end +end + if actualMaxCount <= 0 then return __qmret(nil) end @@ -4632,11 +5337,22 @@ for _, ckQueueName in ipairs(ckQueues) do end local fullQueueKey = keyPrefix .. ckQueueName + local blockedByGates = false local ckConcurrencyKey = fullQueueKey .. ':currentConcurrency' local ckCurrentConcurrency = tonumber(redis.call('SCARD', ckConcurrencyKey) or '0') - if ckCurrentConcurrency < queueConcurrencyLimit then + local perKeyLimit = queueConcurrencyLimit + + if ckCurrentConcurrency >= perKeyLimit then + -- Back a blocked variant off so it cannot pin the bounded candidate window + -- and starve later keys (acute with a zero per-key override, which never + -- self-clears). Acks and nacks rebalance the score back to the oldest + -- message, so the key is eligible again the moment capacity frees. + redis.call('ZADD', ckIndexKey, currentTime + 1000, ckQueueName) + end + + if ckCurrentConcurrency < perKeyLimit then local messages = redis.call('ZRANGEBYSCORE', fullQueueKey, '-inf', tostring(currentTime), 'WITHSCORES', 'LIMIT', 0, 1) if #messages >= 2 then @@ -4651,30 +5367,62 @@ for _, ckQueueName in ipairs(ckQueues) do local ttlExpiresAt = messageData and messageData.ttlExpiresAt if ttlExpiresAt and ttlExpiresAt <= currentTime then - redis.call('ZREM', fullQueueKey, messageId) + local removedExpired = redis.call('ZREM', fullQueueKey, messageId) redis.call('ZREM', envQueueKey, messageId) decrLengthCounter() + if removedExpired == 1 then + __gateQueuedDelta(keyPrefix, messageData, -1) + end if ttlQueueKey and ttlQueueKey ~= '' then local ttlMember = ckQueueName .. '|' .. messageId .. '|' .. (messageData.orgId or '') redis.call('ZADD', ttlQueueKey, ttlExpiresAt, ttlMember) end else - redis.call('ZREM', fullQueueKey, messageId) - redis.call('ZREM', envQueueKey, messageId) - decrLengthCounter() - redis.call('SADD', ckConcurrencyKey, messageId) - redis.call('SADD', envCurrentConcurrencyKey, messageId) + local gatesAllow = true + if gatesEnabled then + gatesAllow = __gatesHaveCapacity(keyPrefix, messageData, messageId, envConcurrencyLimit, messageKeyPrefix) + end + if not gatesAllow then + blockedByGates = true + end - if ttlQueueKey and ttlQueueKey ~= '' and ttlExpiresAt then - local ttlMember = ckQueueName .. '|' .. messageId .. '|' .. (messageData.orgId or '') - redis.call('ZREM', ttlQueueKey, ttlMember) + local alreadyInGroup = false + local totalAllows = true + if totalHeadroom ~= nil then + alreadyInGroup = redis.call('SISMEMBER', groupConcurrencyKey, messageId) == 1 + totalAllows = alreadyInGroup or totalHeadroom > 0 end - table.insert(results, messageId) - table.insert(results, messageScore) - table.insert(results, messagePayload) + if gatesAllow and totalAllows then + local removedFromQueue = redis.call('ZREM', fullQueueKey, messageId) + redis.call('ZREM', envQueueKey, messageId) + decrLengthCounter() + if removedFromQueue == 1 then + __gateQueuedDelta(keyPrefix, messageData, -1) + end + redis.call('SADD', ckConcurrencyKey, messageId) + redis.call('SADD', envCurrentConcurrencyKey, messageId) + if totalConcurrencyEnabled then + redis.call('SADD', groupConcurrencyKey, messageId) + if totalHeadroom ~= nil and not alreadyInGroup then + totalHeadroom = totalHeadroom - 1 + end + end + if gatesEnabled then + __gatesAcquire(keyPrefix, messageData, messageId) + end - dequeuedCount = dequeuedCount + 1 + if ttlQueueKey and ttlQueueKey ~= '' and ttlExpiresAt then + local ttlMember = ckQueueName .. '|' .. messageId .. '|' .. (messageData.orgId or '') + redis.call('ZREM', ttlQueueKey, ttlMember) + end + + table.insert(results, messageId) + table.insert(results, messageScore) + table.insert(results, messagePayload) + + dequeuedCount = dequeuedCount + 1 + end end else redis.call('ZREM', fullQueueKey, messageId) @@ -4682,11 +5430,15 @@ for _, ckQueueName in ipairs(ckQueues) do decrLengthCounter() end - local earliest = redis.call('ZRANGE', fullQueueKey, 0, 0, 'WITHSCORES') - if #earliest == 0 then - redis.call('ZREM', ckIndexKey, ckQueueName) + if blockedByGates then + redis.call('ZADD', ckIndexKey, currentTime + 1000, ckQueueName) else - redis.call('ZADD', ckIndexKey, earliest[2], ckQueueName) + local earliest = redis.call('ZRANGE', fullQueueKey, 0, 0, 'WITHSCORES') + if #earliest == 0 then + redis.call('ZREM', ckIndexKey, ckQueueName) + else + redis.call('ZADD', ckIndexKey, earliest[2], ckQueueName) + end end else local any = redis.call('ZRANGE', fullQueueKey, 0, 0, 'WITHSCORES') @@ -4832,7 +5584,7 @@ return message }); this.redis.defineCommand("acknowledgeMessage", { - numberOfKeys: 9, + numberOfKeys: 10, lua: ` -- Keys: local masterQueueKey = KEYS[1] @@ -4844,19 +5596,27 @@ local queueCurrentDequeuedKey = KEYS[6] local envCurrentDequeuedKey = KEYS[7] local envQueueKey = KEYS[8] local workerQueueKey = KEYS[9] +local groupConcurrencyKey = KEYS[10] -- Args: local messageId = ARGV[1] local messageQueueName = ARGV[2] local messageKeyValue = ARGV[3] local removeFromWorkerQueue = ARGV[4] +local keyPrefix = ARGV[5] +${QUEUE_GATES_LUA_HELPERS} + +local rawPayload = redis.call('GET', messageKey) -- Remove the message from the message key redis.call('DEL', messageKey) -- Remove the message from the queue -redis.call('ZREM', messageQueueKey, messageId) +local removedFromQueueZset = redis.call('ZREM', messageQueueKey, messageId) redis.call('ZREM', envQueueKey, messageId) +if removedFromQueueZset == 1 then + __gateQueuedDeltaRaw(keyPrefix, rawPayload, -1) +end -- Rebalance the parent queues local earliestMessage = redis.call('ZRANGE', messageQueueKey, 0, 0, 'WITHSCORES') @@ -4866,11 +5626,16 @@ else redis.call('ZADD', masterQueueKey, earliestMessage[2], messageQueueName) end --- Update the concurrency keys -redis.call('SREM', queueCurrentConcurrencyKey, messageId) +-- Update the concurrency keys. The groupConcurrency SREM mirrors the base SREM +-- unconditionally (no flag check) so a disabled flag still drains the group set. +local removedFromCurrentConcurrency = redis.call('SREM', queueCurrentConcurrencyKey, messageId) +if removedFromCurrentConcurrency == 1 then + redis.call('SREM', groupConcurrencyKey, messageId) +end redis.call('SREM', envCurrentConcurrencyKey, messageId) redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) +__gatesRelease(keyPrefix, rawPayload, messageId) -- Remove the message from the worker queue if removeFromWorkerQueue == '1' then @@ -4880,7 +5645,7 @@ end }); this.redis.defineCommand("nackMessage", { - numberOfKeys: 8, + numberOfKeys: 9, lua: ` -- Keys: local masterQueueKey = KEYS[1] @@ -4891,24 +5656,35 @@ local envCurrentConcurrencyKey = KEYS[5] local queueCurrentDequeuedKey = KEYS[6] local envCurrentDequeuedKey = KEYS[7] local envQueueKey = KEYS[8] +local groupConcurrencyKey = KEYS[9] -- Args: local messageId = ARGV[1] local messageQueueName = ARGV[2] local messageData = ARGV[3] local messageScore = tonumber(ARGV[4]) +local keyPrefix = ARGV[5] +${QUEUE_GATES_LUA_HELPERS} -- Update the message data redis.call('SET', messageKey, messageData) --- Update the concurrency keys -redis.call('SREM', queueCurrentConcurrencyKey, messageId) +-- Update the concurrency keys. The groupConcurrency SREM mirrors the base SREM +-- unconditionally (no flag check) so a disabled flag still drains the group set. +local removedFromCurrentConcurrency = redis.call('SREM', queueCurrentConcurrencyKey, messageId) +if removedFromCurrentConcurrency == 1 then + redis.call('SREM', groupConcurrencyKey, messageId) +end redis.call('SREM', envCurrentConcurrencyKey, messageId) redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) +__gatesRelease(keyPrefix, messageData, messageId) -- Enqueue the message into the queue -redis.call('ZADD', messageQueueKey, messageScore, messageId) +local added = redis.call('ZADD', messageQueueKey, messageScore, messageId) +if added == 1 then + __gateQueuedDeltaRaw(keyPrefix, messageData, 1) +end redis.call('ZADD', envQueueKey, messageScore, messageId) -- Rebalance the parent queues @@ -4922,7 +5698,7 @@ end }); this.redis.defineCommand("moveToDeadLetterQueue", { - numberOfKeys: 9, + numberOfKeys: 10, lua: ` -- Keys: local masterQueueKey = KEYS[1] @@ -4934,14 +5710,22 @@ local queueCurrentDequeuedKey = KEYS[6] local envCurrentDequeuedKey = KEYS[7] local envQueueKey = KEYS[8] local deadLetterQueueKey = KEYS[9] +local groupConcurrencyKey = KEYS[10] -- Args: local messageId = ARGV[1] local messageQueueName = ARGV[2] +local keyPrefix = ARGV[3] +${QUEUE_GATES_LUA_HELPERS} + +local rawPayload = redis.call('GET', messageKey) -- Remove the message from the queue -redis.call('ZREM', messageQueue, messageId) +local removedFromQueueZset = redis.call('ZREM', messageQueue, messageId) redis.call('ZREM', envQueueKey, messageId) +if removedFromQueueZset == 1 then + __gateQueuedDeltaRaw(keyPrefix, rawPayload, -1) +end -- Rebalance the parent queues local earliestMessage = redis.call('ZRANGE', messageQueue, 0, 0, 'WITHSCORES') @@ -4954,11 +5738,16 @@ end -- Add the message to the dead letter queue redis.call('ZADD', deadLetterQueueKey, tonumber(redis.call('TIME')[1]), messageId) --- Update the concurrency keys -redis.call('SREM', queueCurrentConcurrencyKey, messageId) +-- Update the concurrency keys. The groupConcurrency SREM mirrors the base SREM +-- unconditionally (no flag check) so a disabled flag still drains the group set. +local removedFromCurrentConcurrency = redis.call('SREM', queueCurrentConcurrencyKey, messageId) +if removedFromCurrentConcurrency == 1 then + redis.call('SREM', groupConcurrencyKey, messageId) +end redis.call('SREM', envCurrentConcurrencyKey, messageId) redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) +__gatesRelease(keyPrefix, rawPayload, messageId) `, }); @@ -5153,7 +5942,7 @@ redis.call('SREM', envCurrentDequeuedKey, messageId) // removed something) and runningCounter (when SREM currentDequeued actually // removed something). this.redis.defineCommand("acknowledgeMessageCkTracked", { - numberOfKeys: 12, + numberOfKeys: 13, lua: ` -- Keys: local masterQueueKey = KEYS[1] @@ -5168,6 +5957,7 @@ local workerQueueKey = KEYS[9] local ckIndexKey = KEYS[10] local lengthCounterKey = KEYS[11] local runningCounterKey = KEYS[12] +local groupConcurrencyKey = KEYS[13] -- Args: local messageId = ARGV[1] @@ -5175,6 +5965,10 @@ local messageQueueName = ARGV[2] local messageKeyValue = ARGV[3] local removeFromWorkerQueue = ARGV[4] local ckWildcardName = ARGV[5] +local keyPrefix = ARGV[6] +${QUEUE_GATES_LUA_HELPERS} + +local rawPayload = redis.call('GET', messageKey) local function decrFloored(key) if tonumber(redis.call('GET', key) or '0') > 0 then @@ -5192,6 +5986,7 @@ local removedFromZset = redis.call('ZREM', messageQueueKey, messageId) redis.call('ZREM', envQueueKey, messageId) if removedFromZset == 1 then decrFloored(lengthCounterKey) + __gateQueuedDeltaRaw(keyPrefix, rawPayload, -1) end -- Rebalance CK index @@ -5220,13 +6015,19 @@ end -- Update the concurrency keys. DECR runningCounter only when SREM -- currentDequeued actually removed an entry (the message was in flight). -redis.call('SREM', queueCurrentConcurrencyKey, messageId) +-- The groupConcurrency SREM mirrors the per-CK SREM so the group set drains +-- on every release path. +local removedFromCurrentConcurrency = redis.call('SREM', queueCurrentConcurrencyKey, messageId) +if removedFromCurrentConcurrency == 1 then + redis.call('SREM', groupConcurrencyKey, messageId) +end redis.call('SREM', envCurrentConcurrencyKey, messageId) local removedFromDequeued = redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) if removedFromDequeued == 1 then decrFloored(runningCounterKey) end +__gatesRelease(keyPrefix, rawPayload, messageId) -- Remove the message from the worker queue if removeFromWorkerQueue == '1' then @@ -5239,7 +6040,7 @@ end // runningCounter (floored); ZADD back to the variant zset INCRs // lengthCounter only when ZADD reported a new entry. this.redis.defineCommand("nackMessageCkTracked", { - numberOfKeys: 11, + numberOfKeys: 12, lua: ` -- Keys: local masterQueueKey = KEYS[1] @@ -5253,6 +6054,7 @@ local envQueueKey = KEYS[8] local ckIndexKey = KEYS[9] local lengthCounterKey = KEYS[10] local runningCounterKey = KEYS[11] +local groupConcurrencyKey = KEYS[12] -- Args: local messageId = ARGV[1] @@ -5264,6 +6066,7 @@ local ckWildcardName = ARGV[5] local keyPrefix = ARGV[6] -- TTL (seconds) applied to counter lazy-init SETs local counterTtl = ARGV[7] +${QUEUE_GATES_LUA_HELPERS} local function decrFloored(key) if tonumber(redis.call('GET', key) or '0') > 0 then @@ -5278,13 +6081,18 @@ redis.call('SET', messageKey, messageData) -- so we skip the eager lazy-init here (unlike releaseConcurrencyTracked, which -- mirrors the same DECR pattern with init). A post-TTL nack's floored DECR -- no-ops; the next dequeueMessageFromKeyTracked reseeds from current state. -redis.call('SREM', queueCurrentConcurrencyKey, messageId) +-- The groupConcurrency SREM mirrors the per-CK SREM. +local removedFromCurrentConcurrency = redis.call('SREM', queueCurrentConcurrencyKey, messageId) +if removedFromCurrentConcurrency == 1 then + redis.call('SREM', groupConcurrencyKey, messageId) +end redis.call('SREM', envCurrentConcurrencyKey, messageId) local removedFromDequeued = redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) if removedFromDequeued == 1 then decrFloored(runningCounterKey) end +__gatesRelease(keyPrefix, messageData, messageId) -- Lazy-init lengthCounter if missing (e.g. expired via 24h TTL). nack re-queues a -- message, which means lengthCounter must be present before we INCR. Without this, @@ -5302,6 +6110,9 @@ end -- Enqueue the message back into the CK-specific queue. INCR lengthCounter only if -- it's a new entry (ZADD returns 1). local added = redis.call('ZADD', messageQueueKey, messageScore, messageId) +if added == 1 then + __gateQueuedDeltaRaw(keyPrefix, messageData, 1) +end redis.call('ZADD', envQueueKey, messageScore, messageId) if added == 1 then redis.call('INCR', lengthCounterKey) @@ -5334,7 +6145,7 @@ end // Tracked variant: same as moveToDeadLetterQueueCk. ZREM may DECR // lengthCounter (defensive); SREM currentDequeued may DECR runningCounter. this.redis.defineCommand("moveToDeadLetterQueueCkTracked", { - numberOfKeys: 12, + numberOfKeys: 13, lua: ` -- Keys: local masterQueueKey = KEYS[1] @@ -5349,11 +6160,16 @@ local deadLetterQueueKey = KEYS[9] local ckIndexKey = KEYS[10] local lengthCounterKey = KEYS[11] local runningCounterKey = KEYS[12] +local groupConcurrencyKey = KEYS[13] -- Args: local messageId = ARGV[1] local messageQueueName = ARGV[2] local ckWildcardName = ARGV[3] +local keyPrefix = ARGV[4] +${QUEUE_GATES_LUA_HELPERS} + +local rawPayload = redis.call('GET', messageKey) local function decrFloored(key) if tonumber(redis.call('GET', key) or '0') > 0 then @@ -5368,6 +6184,7 @@ local removedFromZset = redis.call('ZREM', messageQueue, messageId) redis.call('ZREM', envQueueKey, messageId) if removedFromZset == 1 then decrFloored(lengthCounterKey) + __gateQueuedDeltaRaw(keyPrefix, rawPayload, -1) end -- Rebalance CK index @@ -5398,34 +6215,48 @@ end redis.call('ZADD', deadLetterQueueKey, tonumber(redis.call('TIME')[1]), messageId) -- Update the concurrency keys. DECR runningCounter only when SREM --- currentDequeued actually removed an entry. -redis.call('SREM', queueCurrentConcurrencyKey, messageId) +-- currentDequeued actually removed an entry. The groupConcurrency SREM mirrors +-- the per-CK SREM. +local removedFromCurrentConcurrency = redis.call('SREM', queueCurrentConcurrencyKey, messageId) +if removedFromCurrentConcurrency == 1 then + redis.call('SREM', groupConcurrencyKey, messageId) +end redis.call('SREM', envCurrentConcurrencyKey, messageId) local removedFromDequeued = redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) if removedFromDequeued == 1 then decrFloored(runningCounterKey) end +__gatesRelease(keyPrefix, rawPayload, messageId) `, }); this.redis.defineCommand("releaseConcurrency", { - numberOfKeys: 4, + numberOfKeys: 6, lua: ` -- Keys: local queueCurrentConcurrencyKey = KEYS[1] local envCurrentConcurrencyKey = KEYS[2] local queueCurrentDequeuedKey = KEYS[3] local envCurrentDequeuedKey = KEYS[4] +local messageKey = KEYS[5] +local groupConcurrencyKey = KEYS[6] -- Args: local messageId = ARGV[1] +local keyPrefix = ARGV[2] +${QUEUE_GATES_LUA_HELPERS} --- Update the concurrency keys -redis.call('SREM', queueCurrentConcurrencyKey, messageId) +-- Update the concurrency keys. The groupConcurrency SREM mirrors the base SREM +-- unconditionally (no flag check) so a disabled flag still drains the group set. +local removedFromCurrentConcurrency = redis.call('SREM', queueCurrentConcurrencyKey, messageId) +if removedFromCurrentConcurrency == 1 then + redis.call('SREM', groupConcurrencyKey, messageId) +end redis.call('SREM', envCurrentConcurrencyKey, messageId) redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) +__gatesRelease(keyPrefix, redis.call('GET', messageKey), messageId) `, }); @@ -5434,7 +6265,7 @@ redis.call('SREM', envCurrentDequeuedKey, messageId) // something. Caller should only invoke this variant for CK queues — non-CK // queues should keep calling releaseConcurrency. this.redis.defineCommand("releaseConcurrencyTracked", { - numberOfKeys: 6, + numberOfKeys: 8, lua: ` -- Keys: local queueCurrentConcurrencyKey = KEYS[1] @@ -5443,12 +6274,15 @@ local queueCurrentDequeuedKey = KEYS[3] local envCurrentDequeuedKey = KEYS[4] local runningCounterKey = KEYS[5] local ckIndexKey = KEYS[6] +local groupConcurrencyKey = KEYS[7] +local messageKey = KEYS[8] -- Args: local messageId = ARGV[1] local keyPrefix = ARGV[2] -- TTL (seconds) applied to counter lazy-init SETs local counterTtl = ARGV[3] +${QUEUE_GATES_LUA_HELPERS} -- Lazy-init runningCounter if missing (e.g. expired via 24h TTL). Runs BEFORE -- the SREM so the seed captures pre-release state; the subsequent DECR accounts @@ -5464,7 +6298,10 @@ if redis.call('EXISTS', runningCounterKey) == 0 then redis.call('SET', runningCounterKey, total, 'EX', counterTtl) end -redis.call('SREM', queueCurrentConcurrencyKey, messageId) +local removedFromCurrentConcurrency = redis.call('SREM', queueCurrentConcurrencyKey, messageId) +if removedFromCurrentConcurrency == 1 then + redis.call('SREM', groupConcurrencyKey, messageId) +end redis.call('SREM', envCurrentConcurrencyKey, messageId) local removedFromDequeued = redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) @@ -5474,6 +6311,7 @@ if removedFromDequeued == 1 then redis.call('DECR', runningCounterKey) end end +__gatesRelease(keyPrefix, redis.call('GET', messageKey), messageId) `, }); @@ -5557,29 +6395,38 @@ return results }); this.redis.defineCommand("clearMessageFromConcurrencySets", { - numberOfKeys: 4, + numberOfKeys: 6, lua: ` -- Keys: local queueCurrentConcurrencyKey = KEYS[1] local envCurrentConcurrencyKey = KEYS[2] local queueCurrentDequeuedKey = KEYS[3] local envCurrentDequeuedKey = KEYS[4] +local messageKey = KEYS[5] +local groupConcurrencyKey = KEYS[6] -- Args: local messageId = ARGV[1] +local keyPrefix = ARGV[2] +${QUEUE_GATES_LUA_HELPERS} --- Update the concurrency keys -redis.call('SREM', queueCurrentConcurrencyKey, messageId) +-- Update the concurrency keys. The groupConcurrency SREM mirrors the base SREM +-- unconditionally (no flag check) so a disabled flag still drains the group set. +local removedFromCurrentConcurrency = redis.call('SREM', queueCurrentConcurrencyKey, messageId) +if removedFromCurrentConcurrency == 1 then + redis.call('SREM', groupConcurrencyKey, messageId) +end redis.call('SREM', envCurrentConcurrencyKey, messageId) redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) +__gatesRelease(keyPrefix, redis.call('GET', messageKey), messageId) `, }); // Tracked variant of clearMessageFromConcurrencySets — see releaseConcurrencyTracked // for the contract. Only invoke for CK queues. this.redis.defineCommand("clearMessageFromConcurrencySetsTracked", { - numberOfKeys: 6, + numberOfKeys: 8, lua: ` -- Keys: local queueCurrentConcurrencyKey = KEYS[1] @@ -5588,12 +6435,15 @@ local queueCurrentDequeuedKey = KEYS[3] local envCurrentDequeuedKey = KEYS[4] local runningCounterKey = KEYS[5] local ckIndexKey = KEYS[6] +local groupConcurrencyKey = KEYS[7] +local messageKey = KEYS[8] -- Args: local messageId = ARGV[1] local keyPrefix = ARGV[2] -- TTL (seconds) applied to counter lazy-init SETs local counterTtl = ARGV[3] +${QUEUE_GATES_LUA_HELPERS} -- Lazy-init runningCounter if missing — see releaseConcurrencyTracked for rationale. if redis.call('EXISTS', runningCounterKey) == 0 then @@ -5605,7 +6455,10 @@ if redis.call('EXISTS', runningCounterKey) == 0 then redis.call('SET', runningCounterKey, total, 'EX', counterTtl) end -redis.call('SREM', queueCurrentConcurrencyKey, messageId) +local removedFromCurrentConcurrency = redis.call('SREM', queueCurrentConcurrencyKey, messageId) +if removedFromCurrentConcurrency == 1 then + redis.call('SREM', groupConcurrencyKey, messageId) +end redis.call('SREM', envCurrentConcurrencyKey, messageId) local removedFromDequeued = redis.call('SREM', queueCurrentDequeuedKey, messageId) redis.call('SREM', envCurrentDequeuedKey, messageId) @@ -5615,6 +6468,7 @@ if removedFromDequeued == 1 then redis.call('DECR', runningCounterKey) end end +__gatesRelease(keyPrefix, redis.call('GET', messageKey), messageId) `, }); } @@ -5644,6 +6498,8 @@ declare module "@internal/redis" { queueConcurrencyLimitKey: string, envConcurrencyLimitKey: string, envConcurrencyLimitBurstFactorKey: string, + groupConcurrencyKey: string, + totalConcurrencyLimitKey: string, //args queueName: string, messageId: string, @@ -5654,6 +6510,9 @@ declare module "@internal/redis" { defaultEnvConcurrencyBurstFactor: string, currentTime: string, enableFastPath: string, + keyPrefix: string, + gatesEnabled: string, + totalConcurrencyEnabled: string, metricsEnabled: string, callback?: Callback<[number, number[] | null]> ): Result<[number, number[] | null], Context>; @@ -5673,6 +6532,8 @@ declare module "@internal/redis" { queueConcurrencyLimitKey: string, envConcurrencyLimitKey: string, envConcurrencyLimitBurstFactorKey: string, + groupConcurrencyKey: string, + totalConcurrencyLimitKey: string, //args queueName: string, messageId: string, @@ -5685,6 +6546,9 @@ declare module "@internal/redis" { defaultEnvConcurrencyBurstFactor: string, currentTime: string, enableFastPath: string, + keyPrefix: string, + gatesEnabled: string, + totalConcurrencyEnabled: string, metricsEnabled: string, callback?: Callback<[number, number[] | null]> ): Result<[number, number[] | null], Context>; @@ -5715,6 +6579,8 @@ declare module "@internal/redis" { envQueueKey: string, masterQueueKey: string, ttlQueueKey: string, + groupConcurrencyKey: string, + totalConcurrencyLimitKey: string, //args childQueueName: string, currentTime: string, @@ -5722,6 +6588,8 @@ declare module "@internal/redis" { defaultEnvConcurrencyBurstFactor: string, keyPrefix: string, maxCount: string, + gatesEnabled: string, + totalConcurrencyEnabled: string, metricsEnabled: string, callback?: Callback<[string[] | null, number[] | null]> ): Result<[string[] | null, number[] | null], Context>; @@ -5750,11 +6618,13 @@ declare module "@internal/redis" { envCurrentDequeuedKey: string, envQueueKey: string, workerQueueKey: string, + groupConcurrencyKey: string, // args messageId: string, messageQueueName: string, messageKeyValue: string, removeFromWorkerQueue: string, + keyPrefix: string, callback?: Callback ): Result; @@ -5764,8 +6634,11 @@ declare module "@internal/redis" { envCurrentConcurrencyKey: string, queueCurrentDequeuedKey: string, envCurrentDequeuedKey: string, + messageKey: string, + groupConcurrencyKey: string, // args messageId: string, + keyPrefix: string, callback?: Callback ): Result; @@ -5779,11 +6652,13 @@ declare module "@internal/redis" { queueCurrentDequeuedKey: string, envCurrentDequeuedKey: string, envQueueKey: string, + groupConcurrencyKey: string, // args messageId: string, messageQueueName: string, messageData: string, messageScore: string, + keyPrefix: string, callback?: Callback ): Result; @@ -5798,9 +6673,11 @@ declare module "@internal/redis" { envCurrentDequeuedKey: string, envQueueKey: string, deadLetterQueueKey: string, + groupConcurrencyKey: string, // args messageId: string, messageQueueName: string, + keyPrefix: string, callback?: Callback ): Result; @@ -5810,8 +6687,11 @@ declare module "@internal/redis" { envCurrentConcurrencyKey: string, queueCurrentDequeuedKey: string, envCurrentDequeuedKey: string, + messageKey: string, + groupConcurrencyKey: string, // args messageId: string, + keyPrefix: string, callback?: Callback ): Result; @@ -5999,6 +6879,8 @@ declare module "@internal/redis" { envConcurrencyLimitBurstFactorKey: string, lengthCounterKey: string, baseQueueKey: string, + groupConcurrencyKey: string, + totalConcurrencyLimitKey: string, queueName: string, messageId: string, messageData: string, @@ -6011,6 +6893,8 @@ declare module "@internal/redis" { enableFastPath: string, keyPrefix: string, counterTtl: string, + totalConcurrencyEnabled: string, + gatesEnabled: string, metricsEnabled: string, callback?: Callback<[number, number[] | null]> ): Result<[number, number[] | null], Context>; @@ -6032,6 +6916,8 @@ declare module "@internal/redis" { envConcurrencyLimitBurstFactorKey: string, lengthCounterKey: string, baseQueueKey: string, + groupConcurrencyKey: string, + totalConcurrencyLimitKey: string, queueName: string, messageId: string, messageData: string, @@ -6046,6 +6932,8 @@ declare module "@internal/redis" { enableFastPath: string, keyPrefix: string, counterTtl: string, + totalConcurrencyEnabled: string, + gatesEnabled: string, metricsEnabled: string, callback?: Callback<[number, number[] | null]> ): Result<[number, number[] | null], Context>; @@ -6062,12 +6950,16 @@ declare module "@internal/redis" { ttlQueueKey: string, lengthCounterKey: string, runningCounterKey: string, + groupConcurrencyKey: string, + totalConcurrencyLimitKey: string, ckWildcardName: string, currentTime: string, defaultEnvConcurrencyLimit: string, defaultEnvConcurrencyBurstFactor: string, keyPrefix: string, maxCount: string, + totalConcurrencyEnabled: string, + gatesEnabled: string, metricsEnabled: string, callback?: Callback<[string[] | null, number[] | null]> ): Result<[string[] | null, number[] | null], Context>; @@ -6092,11 +6984,13 @@ declare module "@internal/redis" { ckIndexKey: string, lengthCounterKey: string, runningCounterKey: string, + groupConcurrencyKey: string, messageId: string, messageQueueName: string, messageKeyValue: string, removeFromWorkerQueue: string, ckWildcardName: string, + keyPrefix: string, callback?: Callback ): Result; @@ -6112,6 +7006,7 @@ declare module "@internal/redis" { ckIndexKey: string, lengthCounterKey: string, runningCounterKey: string, + groupConcurrencyKey: string, messageId: string, messageQueueName: string, messageData: string, @@ -6135,9 +7030,11 @@ declare module "@internal/redis" { ckIndexKey: string, lengthCounterKey: string, runningCounterKey: string, + groupConcurrencyKey: string, messageId: string, messageQueueName: string, ckWildcardName: string, + keyPrefix: string, callback?: Callback ): Result; @@ -6160,6 +7057,8 @@ declare module "@internal/redis" { envCurrentDequeuedKey: string, runningCounterKey: string, ckIndexKey: string, + groupConcurrencyKey: string, + messageKey: string, messageId: string, keyPrefix: string, counterTtl: string, @@ -6173,6 +7072,8 @@ declare module "@internal/redis" { envCurrentDequeuedKey: string, runningCounterKey: string, ckIndexKey: string, + groupConcurrencyKey: string, + messageKey: string, messageId: string, keyPrefix: string, counterTtl: string, diff --git a/internal-packages/run-engine/src/run-queue/keyProducer.ts b/internal-packages/run-engine/src/run-queue/keyProducer.ts index 0609b3d719b..ff7a51f91f3 100644 --- a/internal-packages/run-engine/src/run-queue/keyProducer.ts +++ b/internal-packages/run-engine/src/run-queue/keyProducer.ts @@ -24,6 +24,9 @@ const constants = { CK_INDEX_PART: "ckIndex", LENGTH_COUNTER_PART: "lengthCounter", RUNNING_COUNTER_PART: "runningCounter", + GROUP_CONCURRENCY_PART: "groupConcurrency", + TOTAL_CONCURRENCY_LIMIT_PART: "totalConcurrency", + GATE_QUEUED_COUNTER_PART: "gateQueuedCounter", } as const; export class RunQueueFullKeyProducer implements RunQueueKeyProducer { @@ -338,6 +341,42 @@ export class RunQueueFullKeyProducer implements RunQueueKeyProducer { return `${this.baseQueueKeyFromQueue(queue)}:${constants.RUNNING_COUNTER_PART}`; } + /** + * SET of in-flight messageIds across ALL concurrency-key variants of a base queue. + * SCARD of this set is the queue's total running count, gated by the total + * concurrency limit. Lives at the base queue so every ck variant shares it. + */ + queueGroupConcurrencyKey(env: RunQueueKeyProducerEnvironment, queue: string): string { + return `${this.queueKey(env, queue)}:${constants.GROUP_CONCURRENCY_PART}`; + } + + queueGroupConcurrencyKeyFromQueue(queue: string): string { + return `${this.baseQueueKeyFromQueue(queue)}:${constants.GROUP_CONCURRENCY_PART}`; + } + + /** + * Counter of queued runs holding this queue as a gate: runs that are not + * executing, are queued, and must clear this gate to execute. Maintained + * exactly like the CK length counter: incremented per gate on enqueue, + * decremented on admit and on every queued-removal path. + */ + gateQueuedCounterKey(env: RunQueueKeyProducerEnvironment, queue: string): string { + return `${this.queueKey(env, queue)}:${constants.GATE_QUEUED_COUNTER_PART}`; + } + + /** + * String key holding the queue's total concurrency limit (the cap across all + * concurrency-key variants). Absent = no total cap. Readers clamp to the + * environment limit; the raw requested value is what's stored. + */ + queueTotalConcurrencyLimitKey(env: RunQueueKeyProducerEnvironment, queue: string): string { + return `${this.queueKey(env, queue)}:${constants.TOTAL_CONCURRENCY_LIMIT_PART}`; + } + + queueTotalConcurrencyLimitKeyFromQueue(queue: string): string { + return `${this.baseQueueKeyFromQueue(queue)}:${constants.TOTAL_CONCURRENCY_LIMIT_PART}`; + } + isCkWildcard(queue: string): boolean { return queue.endsWith(":ck:*"); } diff --git a/internal-packages/run-engine/src/run-queue/metrics.test.ts b/internal-packages/run-engine/src/run-queue/metrics.test.ts index ebfc295470e..c239b4aa937 100644 --- a/internal-packages/run-engine/src/run-queue/metrics.test.ts +++ b/internal-packages/run-engine/src/run-queue/metrics.test.ts @@ -24,6 +24,21 @@ const authenticatedEnvDev = { organization: { id: "o1234" }, }; +// A dead Redis leaves waitUntilReady() pending forever (the client retries +// indefinitely), which would burn the whole test timeout with no diagnostic. +// The abort releases the losing timer promptly so it cannot hold an event +// loop open for the remaining 15s after a fast ready. +async function emitterReady(emitter: MetricsStreamEmitter) { + const abort = new AbortController(); + const timedOut = setTimeout(15_000, "timeout", { signal: abort.signal }).catch(() => "aborted"); + const winner = await Promise.race([emitter.waitUntilReady().then(() => "ready"), timedOut]); + abort.abort(); + if (winner === "timeout") { + void emitter.close().catch(() => {}); + throw new Error("metrics emitter Redis connection never became ready"); + } +} + async function readAllEntries( redisOptions: { host: string; @@ -81,6 +96,7 @@ describe("RunQueue queue-metrics emission", () => { definition, flag: { enabled: () => true }, }); + await emitterReady(emitter); const queue = new RunQueue({ name: "rq", @@ -123,7 +139,10 @@ describe("RunQueue queue-metrics emission", () => { const entries = await waitForEntries(redis, definition, (es) => { const seen = es.map((e) => e.fields.op); - return ["enqueue", "gauge", "started", "ack"].every((o) => seen.includes(o)); + if (!["enqueue", "gauge", "started", "ack"].every((o) => seen.includes(o))) return false; + return es.some( + (e) => e.fields.op === "gauge" && e.fields.cc === "1" && e.fields.ql === "0" + ); }); const ops = entries.map((e) => e.fields.op); expect(ops).toContain("enqueue"); @@ -137,9 +156,23 @@ describe("RunQueue queue-metrics emission", () => { for (const f of ["ql", "cc", "lim", "eql", "ec", "elim", "thr"]) { expect(gauge!.fields[f]).toBeDefined(); } - // Non-CK scripts keep the 7-field gauge (no CK-health tail). - expect(gauge!.fields.ckq).toBeUndefined(); - expect(gauge!.fields.ckw).toBeUndefined(); + /** + * Non-CK scripts emit the full gauge tail too: zeroed CK-health fields (a base + * queue has no CK variants) followed by the total running/limit pair, so a + * keyless queue with a total limit still charts total concurrency. + */ + expect(gauge!.fields.ckq).toBe("0"); + expect(gauge!.fields.ckw).toBe("0"); + expect(gauge!.fields.tcc).toBeDefined(); + expect(gauge!.fields.tlim).toBeDefined(); + + // Pins the dequeue script's sample-at-return wrapper: only the dequeue emits the + // post-admission reading (running 1, queued 0); the enqueue gauge sees the inverse. + const dequeueGauge = entries.find( + (e) => e.fields.op === "gauge" && e.fields.cc === "1" && e.fields.ql === "0" + ); + assertGauge(dequeueGauge); + expect(dequeueGauge!.fields.q).toContain("task/my-task"); // The first counter emission also seeds a cum=0 baseline (no wait); the real reading // carries wait. Pick the reading (cum > 0). @@ -172,6 +205,7 @@ describe("RunQueue queue-metrics emission", () => { definition, flag: { enabled: () => true }, }); + await emitterReady(emitter); const queue = new RunQueue({ name: "rq", tracer: trace.getTracer("rq"), @@ -244,6 +278,7 @@ describe("RunQueue queue-metrics emission", () => { maxLen: 1000, }; const emitter = new MetricsStreamEmitter({ redis, definition, flag: { enabled: () => true } }); + await emitterReady(emitter); const queue = new RunQueue({ name: "rq", tracer: trace.getTracer("rq"), @@ -283,14 +318,13 @@ describe("RunQueue queue-metrics emission", () => { expect(dequeued?.messageId).toBe(message.runId); const entries = await waitForEntries(redis, definition, (es) => - es.some( - (e) => e.fields.op === "gauge" && e.fields.q.includes(":ck:") && e.fields.thr === "0" - ) + es.some((e) => e.fields.op === "gauge" && e.fields.q.includes(":ck:*")) ); const gauges = entries.filter((e) => e.fields.op === "gauge"); expect(gauges.length).toBeGreaterThan(0); - // The aggregate CK dequeue gauge targets the CK wildcard and never sets thr. - const aggregate = gauges.find((e) => e.fields.q.includes(":ck:") && e.fields.thr === "0"); + // The aggregate gauge targets the CK wildcard and only the CK dequeue script emits + // it, so this pins that script's sample-at-return wrapper. + const aggregate = gauges.find((e) => e.fields.q.includes(":ck:*")); assertGauge(aggregate); expect(Number(aggregate!.fields.ql)).toBeGreaterThanOrEqual(0); expect(Number(aggregate!.fields.cc)).toBeGreaterThanOrEqual(0); @@ -344,6 +378,7 @@ describe("RunQueue queue-metrics emission", () => { flag: { enabled: () => true }, gaugeSampleRate: 0, }); + await emitterReady(emitter); const queue = new RunQueue({ name: "rq", tracer: trace.getTracer("rq"), diff --git a/internal-packages/run-engine/src/run-queue/tests/queueGates.test.ts b/internal-packages/run-engine/src/run-queue/tests/queueGates.test.ts new file mode 100644 index 00000000000..272523b9595 --- /dev/null +++ b/internal-packages/run-engine/src/run-queue/tests/queueGates.test.ts @@ -0,0 +1,491 @@ +import { redisTest } from "@internal/testcontainers"; +import { trace } from "@internal/tracing"; +import { setTimeout } from "node:timers/promises"; +import { describe } from "vitest"; +import { FairQueueSelectionStrategy } from "../fairQueueSelectionStrategy.js"; +import { RunQueue } from "../index.js"; +import { RunQueueFullKeyProducer } from "../keyProducer.js"; +import type { InputPayload } from "../types.js"; +import { Decimal } from "@trigger.dev/database"; + +const testOptions = { + name: "rq", + tracer: trace.getTracer("rq"), + workers: 1, + defaultEnvConcurrency: 25, + retryOptions: { + maxAttempts: 5, + factor: 1.1, + minTimeoutInMs: 100, + maxTimeoutInMs: 1_000, + randomize: true, + }, + keys: new RunQueueFullKeyProducer(), +}; + +const authenticatedEnvDev = { + id: "e1234", + type: "DEVELOPMENT" as const, + maximumConcurrencyLimit: 10, + concurrencyLimitBurstFactor: new Decimal(2.0), + project: { id: "p1234" }, + organization: { id: "o1234" }, +}; + +function createQueue(redisContainer: any, gatesEnabled: boolean) { + return new RunQueue({ + ...testOptions, + gatesEnabled, + queueSelectionStrategy: new FairQueueSelectionStrategy({ + redis: { + keyPrefix: "runqueue:test:", + host: redisContainer.getHost(), + port: redisContainer.getPort(), + }, + keys: testOptions.keys, + }), + redis: { + keyPrefix: "runqueue:test:", + host: redisContainer.getHost(), + port: redisContainer.getPort(), + }, + }); +} + +function makeMessage(overrides: Partial = {}): InputPayload { + return { + runId: "r1", + taskIdentifier: "task/my-task", + orgId: "o1234", + projectId: "p1234", + environmentId: "e1234", + environmentType: "DEVELOPMENT", + queue: "task/my-task", + timestamp: Date.now(), + attempt: 0, + ...overrides, + }; +} + +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(); +} + +async function popWorkerQueue(queue: RunQueue, expected: string): Promise { + const next = await queue.dequeueMessageFromWorkerQueue("consumer-1", "main", { + blockingPop: false, + }); + return next?.messageId === expected; +} + +vi.setConfig({ testTimeout: 60_000 }); + +describe("RunQueue gates", () => { + redisTest("an unkeyed gate caps runs across its holders", async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "shared-gate", 1); + + const now = Date.now(); + for (const i of [0, 1]) { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: `r${i}`, + timestamp: now - 1000 + i, + gates: [{ queue: "shared-gate" }], + }), + workerQueue: "main", + }); + } + + const oneAdmitted = await waitFor( + async () => + (await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "shared-gate")) === 1 + ); + expect(oneAdmitted).toBe(true); + + /** The second run must stay queued while the gate is full. */ + await setTimeout(2000); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "shared-gate")).toBe(1); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + + expect(await popWorkerQueue(queue, "r0")).toBe(true); + await queue.acknowledgeMessage(authenticatedEnvDev.organization.id, "r0"); + + /** Acking r0 frees the gate slot and its home slot; r1 is admitted. */ + const r1Admitted = await waitFor(() => popWorkerQueue(queue, "r1")); + expect(r1Admitted).toBe(true); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "shared-gate")).toBe(1); + } finally { + await queue.quit(); + } + }); + + redisTest( + "a keyed gate caps a tenant across home concurrency keys", + async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "tenant", 1); + + const now = Date.now(); + for (const [i, ck] of ["ck-a", "ck-b"].entries()) { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: `r${i}`, + concurrencyKey: ck, + timestamp: now - 1000 + i, + gates: [{ queue: "tenant", concurrencyKey: "acme" }], + }), + workerQueue: "main", + }); + } + + const oneAdmitted = await waitFor( + async () => + (await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "tenant", "acme")) === 1 + ); + expect(oneAdmitted).toBe(true); + + await setTimeout(2000); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "tenant", "acme")).toBe( + 1 + ); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "tenant")).toBe(1); + + expect(await popWorkerQueue(queue, "r0")).toBe(true); + await queue.acknowledgeMessage(authenticatedEnvDev.organization.id, "r0"); + + const r1Admitted = await waitFor(() => popWorkerQueue(queue, "r1")); + expect(r1Admitted).toBe(true); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "tenant", "acme")).toBe( + 1 + ); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "tenant")).toBe(1); + } finally { + await queue.quit(); + } + } + ); + + redisTest( + "a gate's total limit caps keyed and keyless holders together", + async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueTotalConcurrencyLimits(authenticatedEnvDev, "db-limit", 1); + + const now = Date.now(); + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "r0", + timestamp: now - 1000, + gates: [{ queue: "db-limit" }], + }), + workerQueue: "main", + }); + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "r1", + timestamp: now - 999, + gates: [{ queue: "db-limit", concurrencyKey: "acme" }], + }), + workerQueue: "main", + }); + + const oneAdmitted = await waitFor( + async () => (await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "db-limit")) === 1 + ); + expect(oneAdmitted).toBe(true); + + /** The keyed run must stay queued: the keyless holder occupies the total pool. */ + await setTimeout(2000); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "db-limit")).toBe(1); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + + expect(await popWorkerQueue(queue, "r0")).toBe(true); + await queue.acknowledgeMessage(authenticatedEnvDev.organization.id, "r0"); + + /** Acking the keyless holder frees the total pool; the keyed run is admitted. */ + const r1Admitted = await waitFor(() => popWorkerQueue(queue, "r1")); + expect(r1Admitted).toBe(true); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "db-limit")).toBe(1); + } finally { + await queue.quit(); + } + } + ); + + redisTest( + "counts queued runs per gate and drains the counter on admit and ack", + async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "shared-gate", 1); + + const now = Date.now(); + for (const i of [0, 1]) { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: `r${i}`, + timestamp: now - 1000 + i, + gates: [{ queue: "shared-gate" }], + }), + workerQueue: "main", + }); + } + + const oneAdmitted = await waitFor( + async () => + (await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "shared-gate")) === 1 + ); + expect(oneAdmitted).toBe(true); + + /** One run executes, one waits: the gate's queued counter holds the waiter. */ + await setTimeout(2000); + expect(await queue.gateQueuedCountOfQueue(authenticatedEnvDev, "shared-gate")).toBe(1); + + expect(await popWorkerQueue(queue, "r0")).toBe(true); + await queue.acknowledgeMessage(authenticatedEnvDev.organization.id, "r0"); + + const drained = await waitFor(async () => { + if (!(await popWorkerQueue(queue, "r1"))) return false; + return (await queue.gateQueuedCountOfQueue(authenticatedEnvDev, "shared-gate")) === 0; + }); + expect(drained).toBe(true); + } finally { + await queue.quit(); + } + } + ); + + redisTest("ignores gates and holds no gate slots when disabled", async ({ redisContainer }) => { + const queue = createQueue(redisContainer, false); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "shared-gate", 1); + + const now = Date.now(); + for (const i of [0, 1]) { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: `r${i}`, + timestamp: now - 1000 + i, + gates: [{ queue: "shared-gate" }], + }), + workerQueue: "main", + }); + } + + const bothAdmitted = await waitFor( + async () => (await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")) === 0 + ); + expect(bothAdmitted).toBe(true); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "shared-gate")).toBe(0); + } finally { + await queue.quit(); + } + }); + + redisTest("nacking releases the gate slot", async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "shared-gate", 1); + + const now = Date.now(); + for (const i of [0, 1]) { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: `r${i}`, + timestamp: now - 1000 + i, + gates: [{ queue: "shared-gate" }], + }), + workerQueue: "main", + }); + } + + const r0Admitted = await waitFor(async () => { + if ((await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "shared-gate")) !== 1) { + return false; + } + return popWorkerQueue(queue, "r0"); + }); + expect(r0Admitted).toBe(true); + + await queue.nackMessage({ + orgId: authenticatedEnvDev.organization.id, + messageId: "r0", + retryAt: Date.now() + 120_000, + }); + + /** r0's gate slot must be released so r1 (the only eligible run) is admitted. */ + const r1Admitted = await waitFor(() => popWorkerQueue(queue, "r1")); + expect(r1Admitted).toBe(true); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "shared-gate")).toBe(1); + } finally { + await queue.quit(); + } + }); + + redisTest("reconciles a leaked gate member instead of blocking", async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + const keys = testOptions.keys; + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "shared-gate", 1); + + /** A dead member with no message key occupies the gate. */ + await queue.redis.sadd( + keys.queueCurrentConcurrencyKey(authenticatedEnvDev, "shared-gate"), + "dead-0" + ); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "r0", + timestamp: Date.now() - 1000, + gates: [{ queue: "shared-gate" }], + }), + workerQueue: "main", + }); + + const r0Admitted = await waitFor(() => popWorkerQueue(queue, "r0"), 30_000); + expect(r0Admitted).toBe(true); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "shared-gate")).toBe(1); + } finally { + await queue.quit(); + } + }); + + redisTest( + "a run already holding its own gate slot is never deadlocked by it", + async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + const keys = testOptions.keys; + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "shared-gate", 1); + + /** + * An unmirrored release (an older build's nack) leaves the run's own + * membership behind while the run goes back to waiting in its queue. The + * gate is "full" with the run itself; admission must still let it through. + */ + await queue.redis.sadd( + keys.queueCurrentConcurrencyKey(authenticatedEnvDev, "shared-gate"), + "r0" + ); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "r0", + timestamp: Date.now() - 1000, + gates: [{ queue: "shared-gate" }], + }), + workerQueue: "main", + }); + + const r0Admitted = await waitFor(() => popWorkerQueue(queue, "r0"), 30_000); + expect(r0Admitted).toBe(true); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "shared-gate")).toBe(1); + } finally { + await queue.quit(); + } + } + ); + + redisTest( + "a gate without a key inherits the run's concurrency key", + async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "tenant", 1); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "r0", + concurrencyKey: "acme", + timestamp: Date.now() - 1000, + gates: [{ queue: "tenant" }], + }), + workerQueue: "main", + }); + + const admitted = await waitFor( + async () => + (await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "tenant", "acme")) === 1 + ); + expect(admitted).toBe(true); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "tenant")).toBe(0); + } finally { + await queue.quit(); + } + } + ); + + redisTest("enqueue fast path respects a full gate", async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "shared-gate", 1); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "r0", + timestamp: Date.now() - 1000, + gates: [{ queue: "shared-gate" }], + }), + workerQueue: "main", + enableFastPath: true, + skipDequeueProcessing: true, + }); + + /** The fast path admits synchronously and takes the gate slot. */ + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "shared-gate")).toBe(1); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(0); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "r1", + timestamp: Date.now() - 999, + gates: [{ queue: "shared-gate" }], + }), + workerQueue: "main", + enableFastPath: true, + skipDequeueProcessing: true, + }); + + /** At the gate's limit the fast path must fall back to a normal enqueue. */ + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "shared-gate")).toBe(1); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + } finally { + await queue.quit(); + } + }); +}); diff --git a/internal-packages/run-engine/src/run-queue/tests/totalConcurrency.test.ts b/internal-packages/run-engine/src/run-queue/tests/totalConcurrency.test.ts new file mode 100644 index 00000000000..0ccedd8881b --- /dev/null +++ b/internal-packages/run-engine/src/run-queue/tests/totalConcurrency.test.ts @@ -0,0 +1,589 @@ +import { assertNonNullable, redisTest } from "@internal/testcontainers"; +import { trace } from "@internal/tracing"; +import { setTimeout } from "node:timers/promises"; +import { describe } from "vitest"; +import { FairQueueSelectionStrategy } from "../fairQueueSelectionStrategy.js"; +import { RunQueue } from "../index.js"; +import { RunQueueFullKeyProducer } from "../keyProducer.js"; +import type { InputPayload } from "../types.js"; +import { Decimal } from "@trigger.dev/database"; + +const testOptions = { + name: "rq", + tracer: trace.getTracer("rq"), + workers: 1, + defaultEnvConcurrency: 25, + retryOptions: { + maxAttempts: 5, + factor: 1.1, + minTimeoutInMs: 100, + maxTimeoutInMs: 1_000, + randomize: true, + }, + keys: new RunQueueFullKeyProducer(), +}; + +const authenticatedEnvDev = { + id: "e1234", + type: "DEVELOPMENT" as const, + maximumConcurrencyLimit: 10, + concurrencyLimitBurstFactor: new Decimal(2.0), + project: { id: "p1234" }, + organization: { id: "o1234" }, +}; + +function createQueue(redisContainer: any, totalConcurrencyEnabled: boolean) { + return new RunQueue({ + ...testOptions, + totalConcurrencyEnabled, + queueSelectionStrategy: new FairQueueSelectionStrategy({ + redis: { + keyPrefix: "runqueue:test:", + host: redisContainer.getHost(), + port: redisContainer.getPort(), + }, + keys: testOptions.keys, + }), + redis: { + keyPrefix: "runqueue:test:", + host: redisContainer.getHost(), + port: redisContainer.getPort(), + }, + }); +} + +function makeMessage(overrides: Partial = {}): InputPayload { + return { + runId: "r1", + taskIdentifier: "task/my-task", + orgId: "o1234", + projectId: "p1234", + environmentId: "e1234", + environmentType: "DEVELOPMENT", + queue: "task/my-task", + timestamp: Date.now(), + attempt: 0, + ...overrides, + }; +} + +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(); +} + +vi.setConfig({ testTimeout: 60_000 }); + +describe("RunQueue total concurrency limit", () => { + redisTest( + "caps in-flight runs across concurrency keys at the total limit", + async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueTotalConcurrencyLimits(authenticatedEnvDev, "task/my-task", 2); + + const now = Date.now(); + for (const [i, ck] of ["ck-a", "ck-a", "ck-b", "ck-b"].entries()) { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: `r${i}`, + concurrencyKey: ck, + timestamp: now - 1000 + i, + }), + workerQueue: "main", + }); + } + + const admittedTwo = await waitFor( + async () => + (await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")) === 2 + ); + expect(admittedTwo).toBe(true); + + /** + * The remaining two messages must stay queued: give the master consumers a + * couple of extra polling cycles to prove the gate holds, not just that it + * hadn't caught up yet. + */ + await setTimeout(2000); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(2); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(2); + + const dequeued1 = await queue.dequeueMessageFromWorkerQueue("consumer-1", "main"); + assertNonNullable(dequeued1); + const dequeued2 = await queue.dequeueMessageFromWorkerQueue("consumer-1", "main"); + assertNonNullable(dequeued2); + + await queue.acknowledgeMessage(authenticatedEnvDev.organization.id, dequeued1.messageId); + + const thirdAdmitted = await waitFor(async () => { + const total = await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task"); + const queued = await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task"); + return total === 2 && queued === 1; + }); + expect(thirdAdmitted).toBe(true); + } finally { + await queue.quit(); + } + } + ); + + redisTest( + "still enforces the per-key limit under the total limit", + async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 1); + await queue.updateQueueTotalConcurrencyLimits(authenticatedEnvDev, "task/my-task", 10); + + const now = Date.now(); + for (const i of [0, 1]) { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: `r${i}`, + concurrencyKey: "ck-a", + timestamp: now - 1000 + i, + }), + workerQueue: "main", + }); + } + + const oneAdmitted = await waitFor( + async () => + (await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")) === 1 + ); + expect(oneAdmitted).toBe(true); + + await setTimeout(2000); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + expect( + await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "task/my-task", "ck-a") + ).toBe(1); + } finally { + await queue.quit(); + } + } + ); + + redisTest( + "ignores the stored total limit and maintains no group set when disabled", + async ({ redisContainer }) => { + const queue = createQueue(redisContainer, false); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueTotalConcurrencyLimits(authenticatedEnvDev, "task/my-task", 1); + + const now = Date.now(); + for (const [i, ck] of ["ck-a", "ck-b", "ck-c"].entries()) { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: `r${i}`, + concurrencyKey: ck, + timestamp: now - 1000 + i, + }), + workerQueue: "main", + }); + } + + const allAdmitted = await waitFor( + async () => (await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")) === 0 + ); + expect(allAdmitted).toBe(true); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(0); + } finally { + await queue.quit(); + } + } + ); + + redisTest("the total limit caps keyed and keyless runs together", async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueTotalConcurrencyLimits(authenticatedEnvDev, "task/my-task", 1); + + const now = Date.now(); + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ runId: "r0", timestamp: now - 1000 }), + workerQueue: "main", + }); + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ runId: "r1", concurrencyKey: "ck-a", timestamp: now - 999 }), + workerQueue: "main", + }); + + const oneAdmitted = await waitFor( + async () => (await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")) === 1 + ); + expect(oneAdmitted).toBe(true); + + /** The second run must stay queued: the total pool spans keyed and keyless. */ + await setTimeout(2000); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + + const dequeued = await queue.dequeueMessageFromWorkerQueue("consumer-1", "main"); + assertNonNullable(dequeued); + await queue.acknowledgeMessage(authenticatedEnvDev.organization.id, dequeued.messageId); + + /** Acking the first holder frees the total pool; the other run is admitted. */ + const secondAdmitted = await waitFor(async () => { + const next = await queue.dequeueMessageFromWorkerQueue("consumer-1", "main", { + blockingPop: false, + }); + return next !== undefined && next.messageId !== dequeued.messageId; + }); + expect(secondAdmitted).toBe(true); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + } finally { + await queue.quit(); + } + }); + + redisTest("enqueue fast path respects the total limit", async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueTotalConcurrencyLimits(authenticatedEnvDev, "task/my-task", 1); + + const now = Date.now(); + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ runId: "r0", concurrencyKey: "ck-a", timestamp: now - 1000 }), + workerQueue: "main", + enableFastPath: true, + skipDequeueProcessing: true, + }); + + /** The fast path admits synchronously, so the group slot is taken immediately. */ + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(0); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ runId: "r1", concurrencyKey: "ck-b", timestamp: now - 999 }), + workerQueue: "main", + enableFastPath: true, + skipDequeueProcessing: true, + }); + + /** At the total limit the fast path must fall through to a normal enqueue. */ + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + } finally { + await queue.quit(); + } + }); + + redisTest("nacking releases the total slot", async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueTotalConcurrencyLimits(authenticatedEnvDev, "task/my-task", 1); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ runId: "r0", concurrencyKey: "ck-a", timestamp: Date.now() - 1000 }), + workerQueue: "main", + }); + + const admitted = await waitFor( + async () => (await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")) === 1 + ); + expect(admitted).toBe(true); + + /** A second run on another key waits behind the total limit of 1. */ + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ runId: "r1", concurrencyKey: "ck-b", timestamp: Date.now() - 500 }), + workerQueue: "main", + }); + + await setTimeout(2000); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + + const dequeued = await queue.dequeueMessageFromWorkerQueue("consumer-1", "main"); + assertNonNullable(dequeued); + expect(dequeued.messageId).toBe("r0"); + + /** + * Nack r0 with a far-future retryAt so it cannot immediately reclaim the + * slot. If the nack released r0's group slot, r1 is the only eligible run + * and must be admitted; if the slot leaked, the queue stays blocked and r1 + * never surfaces. + */ + await queue.nackMessage({ + orgId: authenticatedEnvDev.organization.id, + messageId: "r0", + retryAt: Date.now() + 120_000, + }); + + const r1Admitted = await waitFor(async () => { + const next = await queue.dequeueMessageFromWorkerQueue("consumer-1", "main", { + blockingPop: false, + }); + return next?.messageId === "r1"; + }); + expect(r1Admitted).toBe(true); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + } finally { + await queue.quit(); + } + }); + + redisTest( + "reconciles a leaked group member instead of blocking the queue", + async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + const keys = testOptions.keys; + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueTotalConcurrencyLimits(authenticatedEnvDev, "task/my-task", 1); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "r0", + concurrencyKey: "ck-a", + timestamp: Date.now() - 1000, + }), + workerQueue: "main", + }); + + const admitted = await waitFor( + async () => + (await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")) === 1 + ); + expect(admitted).toBe(true); + + const dequeued = await queue.dequeueMessageFromWorkerQueue("consumer-1", "main"); + assertNonNullable(dequeued); + expect(dequeued.messageId).toBe("r0"); + + /** + * Simulate a terminal release from a build without the group mirror: the + * message key is deleted and the per-key and env sets are cleared, but the + * group member is left behind. + */ + await queue.redis.del(keys.messageKey(authenticatedEnvDev.organization.id, "r0")); + await queue.redis.srem( + keys.queueCurrentConcurrencyKey(authenticatedEnvDev, "task/my-task", "ck-a"), + "r0" + ); + await queue.redis.srem(keys.envCurrentConcurrencyKey(authenticatedEnvDev), "r0"); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + + /** The next run must still be admitted: the gate prunes the dead member. */ + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "r1", + concurrencyKey: "ck-b", + timestamp: Date.now() - 500, + }), + workerQueue: "main", + }); + + const r1Admitted = await waitFor(async () => { + const next = await queue.dequeueMessageFromWorkerQueue("consumer-1", "main", { + blockingPop: false, + }); + return next?.messageId === "r1"; + }); + expect(r1Admitted).toBe(true); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + } finally { + await queue.quit(); + } + } + ); + + redisTest( + "reconciles a member whose run went back to waiting in a queue", + async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + const keys = testOptions.keys; + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueTotalConcurrencyLimits(authenticatedEnvDev, "task/my-task", 1); + + /** + * A run on ANOTHER queue whose message exists and is queued there, leaked + * into this queue's group set by an unmirrored release. It can never be a + * dequeue candidate here, so only the reconcile can free the slot: queued + * and in-flight are mutually exclusive, so a queued member holds nothing. + */ + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "x0", + queue: "other-queue", + concurrencyKey: "ck-x", + timestamp: Date.now() + 3_600_000, + }), + workerQueue: "main", + }); + await queue.redis.sadd( + keys.queueGroupConcurrencyKey(authenticatedEnvDev, "task/my-task"), + "x0" + ); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "r0", + concurrencyKey: "ck-a", + timestamp: Date.now() - 1000, + }), + workerQueue: "main", + }); + + const r0Admitted = await waitFor(async () => { + await queue.redis.del( + `${keys.queueGroupConcurrencyKey(authenticatedEnvDev, "task/my-task")}:reconcileLock` + ); + const next = await queue.dequeueMessageFromWorkerQueue("consumer-1", "main", { + blockingPop: false, + }); + return next?.messageId === "r0"; + }, 30_000); + expect(r0Admitted).toBe(true); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + } finally { + await queue.quit(); + } + } + ); + + redisTest( + "reconciles a member whose run dead-lettered or suspended without the mirror", + async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + const keys = testOptions.keys; + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueTotalConcurrencyLimits(authenticatedEnvDev, "task/my-task", 1); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "r0", + concurrencyKey: "ck-a", + timestamp: Date.now() - 1000, + }), + workerQueue: "main", + }); + + const admitted = await waitFor( + async () => + (await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")) === 1 + ); + expect(admitted).toBe(true); + + const dequeued = await queue.dequeueMessageFromWorkerQueue("consumer-1", "main"); + assertNonNullable(dequeued); + expect(dequeued.messageId).toBe("r0"); + + /** + * Simulate a dead-letter or checkpoint release from a build without the + * group mirror: the per-key and env sets are cleared but the MESSAGE KEY + * SURVIVES (dead-letter keeps it for redrive; suspension keeps it until + * the terminal ack) and the run sits in no queue zset. Only the + * home-currentConcurrency reconcile rule can prune this member. + */ + await queue.redis.srem( + keys.queueCurrentConcurrencyKey(authenticatedEnvDev, "task/my-task", "ck-a"), + "r0" + ); + await queue.redis.srem(keys.envCurrentConcurrencyKey(authenticatedEnvDev), "r0"); + expect( + await queue.redis.exists(keys.messageKey(authenticatedEnvDev.organization.id, "r0")) + ).toBe(1); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "r1", + concurrencyKey: "ck-b", + timestamp: Date.now() - 500, + }), + workerQueue: "main", + }); + + const r1Admitted = await waitFor(async () => { + await queue.redis.del( + `${keys.queueGroupConcurrencyKey(authenticatedEnvDev, "task/my-task")}:reconcileLock` + ); + const next = await queue.dequeueMessageFromWorkerQueue("consumer-1", "main", { + blockingPop: false, + }); + return next?.messageId === "r1"; + }, 30_000); + expect(r1Admitted).toBe(true); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + } finally { + await queue.quit(); + } + } + ); + + redisTest( + "reconciles a large leaked backlog across bounded passes", + async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + const keys = testOptions.keys; + const groupKey = keys.queueGroupConcurrencyKey(authenticatedEnvDev, "task/my-task"); + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueTotalConcurrencyLimits(authenticatedEnvDev, "task/my-task", 1); + + /** 1,200 dead members: more than one SSCAN batch, none with a message key. */ + const dead = Array.from({ length: 1200 }, (_, i) => `dead-${i}`); + await queue.redis.sadd(groupKey, ...dead); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1200); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "r0", + concurrencyKey: "ck-a", + timestamp: Date.now() - 1000, + }), + workerQueue: "main", + }); + + /** + * Each dequeue attempt reconciles at most one SSCAN batch behind a 10s + * lock; dropping the lock between polls lets the passes run back to + * back instead of waiting out the interval. + */ + const r0Admitted = await waitFor(async () => { + await queue.redis.del(`${groupKey}:reconcileLock`); + const next = await queue.dequeueMessageFromWorkerQueue("consumer-1", "main", { + blockingPop: false, + }); + return next?.messageId === "r0"; + }, 30_000); + expect(r0Admitted).toBe(true); + expect(await queue.totalConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + } finally { + await queue.quit(); + } + } + ); +}); diff --git a/internal-packages/run-engine/src/run-queue/types.ts b/internal-packages/run-engine/src/run-queue/types.ts index eb8c0f74b70..0a9fa4a8250 100644 --- a/internal-packages/run-engine/src/run-queue/types.ts +++ b/internal-packages/run-engine/src/run-queue/types.ts @@ -2,6 +2,19 @@ import { z } from "zod"; import { RuntimeEnvironmentType } from "@trigger.dev/database"; import type { MinimalAuthenticatedEnvironment } from "../shared/index.js"; +/** + * A gate is another declared queue this run must also hold a concurrency slot in + * while it executes. The run waits in its own queue; each gate contributes an + * extra admit condition (the gate queue's per-key limit, and its total limit when + * 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. + */ +export const QueueGate = z.object({ + queue: z.string().min(1).max(128), + concurrencyKey: z.string().min(1).max(128).optional(), +}); +export type QueueGate = z.infer; + export const InputPayload = z.object({ runId: z.string(), /** Deprecated: not read on the V2 dequeue path; will stop being written in a follow-up. Optional to keep new readers compatible with old payloads that still include it, and vice versa. */ @@ -23,6 +36,9 @@ export const InputPayload = z.object({ // residency so a poll-lagging consumer honors its true residency. Opaque (`z.unknown`) so it // survives serialize while an old consumer strips it — mixed-version safe. Validated at consumption. snapshotRoute: z.unknown().optional(), + /** Additional queues this run must also hold a slot in while executing. At most + * four: three requested gates plus the task's anonymous inline-limit gate. */ + gates: QueueGate.array().max(4).optional(), }); export type InputPayload = z.infer; @@ -97,6 +113,11 @@ export interface RunQueueKeyProducer { queue: string, concurrencyKey?: string ): string; + queueGroupConcurrencyKey(env: RunQueueKeyProducerEnvironment, queue: string): string; + queueGroupConcurrencyKeyFromQueue(queue: string): string; + queueTotalConcurrencyLimitKey(env: RunQueueKeyProducerEnvironment, queue: string): string; + queueTotalConcurrencyLimitKeyFromQueue(queue: string): string; + gateQueuedCounterKey(env: RunQueueKeyProducerEnvironment, queue: string): string; //env oncurrency envCurrentConcurrencyKey(env: EnvDescriptor): 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..3f93c471e92 --- /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 IF NOT EXISTS "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 277bbd7015f..c92cd67a5a8 100644 --- a/internal-packages/run-store/src/types.ts +++ b/internal-packages/run-store/src/types.ts @@ -195,6 +195,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/dev/devSupervisor.ts b/packages/cli-v3/src/dev/devSupervisor.ts index d1cd43b31ad..72a0ba42d01 100644 --- a/packages/cli-v3/src/dev/devSupervisor.ts +++ b/packages/cli-v3/src/dev/devSupervisor.ts @@ -394,6 +394,7 @@ class DevSupervisor implements WorkerRuntime { tasks: backgroundWorker.manifest.tasks, prompts: backgroundWorker.manifest.prompts, queues: backgroundWorker.manifest.queues, + concurrencyLimits: backgroundWorker.manifest.concurrencyLimits, contentHash: manifest.contentHash, sourceFiles, runtime: backgroundWorker.manifest.runtime, 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-controller.ts b/packages/cli-v3/src/entryPoints/managed-index-controller.ts index e4c631d9b51..57569dbe792 100644 --- a/packages/cli-v3/src/entryPoints/managed-index-controller.ts +++ b/packages/cli-v3/src/entryPoints/managed-index-controller.ts @@ -104,6 +104,7 @@ async function indexDeployment({ tasks: workerManifest.tasks, prompts: workerManifest.prompts, queues: workerManifest.queues, + concurrencyLimits: workerManifest.concurrencyLimits, sourceFiles, runtime: workerManifest.runtime, runtimeVersion: workerManifest.runtimeVersion, 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/apiClient/index.ts b/packages/core/src/v3/apiClient/index.ts index 58d074e61b4..852f0f713c4 100644 --- a/packages/core/src/v3/apiClient/index.ts +++ b/packages/core/src/v3/apiClient/index.ts @@ -21,6 +21,9 @@ import { type CreateWaitpointTokenRequestBody, type EndAndContinueSessionRequestBody, type ListQueueOptions, + type ListConcurrencyLimitOptions, + ConcurrencyLimitItem, + type OverrideConcurrencyLimitRequestBody, type ListScheduleOptions, type ListSessionsOptions, type PromotePromptVersionRequestBody, @@ -1760,6 +1763,64 @@ export class ApiClient { ); } + listConcurrencyLimits(options?: ListConcurrencyLimitOptions, requestOptions?: ZodFetchOptions) { + return zodfetchOffsetLimitPage( + ConcurrencyLimitItem, + `${this.baseUrl}/api/v1/concurrency-limits`, + { + page: options?.page, + limit: options?.perPage, + }, + { + method: "GET", + headers: this.#getHeaders(false), + }, + mergeRequestOptions(this.defaultRequestOptions, requestOptions) + ); + } + + retrieveConcurrencyLimit(name: string, requestOptions?: ZodFetchOptions) { + return zodfetch( + ConcurrencyLimitItem, + `${this.baseUrl}/api/v1/concurrency-limits/${encodeURIComponent(name)}`, + { + method: "GET", + headers: this.#getHeaders(false), + }, + mergeRequestOptions(this.defaultRequestOptions, requestOptions) + ); + } + + overrideConcurrencyLimit( + name: string, + override: OverrideConcurrencyLimitRequestBody, + requestOptions?: ZodFetchOptions + ) { + return zodfetch( + ConcurrencyLimitItem, + `${this.baseUrl}/api/v1/concurrency-limits/${encodeURIComponent(name)}/override`, + { + method: "POST", + headers: this.#getHeaders(false), + body: JSON.stringify(override), + }, + mergeRequestOptions(this.defaultRequestOptions, requestOptions) + ); + } + + resetConcurrencyLimit(name: string, requestOptions?: ZodFetchOptions) { + return zodfetch( + ConcurrencyLimitItem, + `${this.baseUrl}/api/v1/concurrency-limits/${encodeURIComponent(name)}/reset`, + { + method: "POST", + headers: this.#getHeaders(false), + body: JSON.stringify({}), + }, + mergeRequestOptions(this.defaultRequestOptions, requestOptions) + ); + } + subscribeToRun( runId: string, options?: { 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 731032613d2..20b28c72156 100644 --- a/packages/core/src/v3/schemas/api.ts +++ b/packages/core/src/v3/schemas/api.ts @@ -282,6 +282,20 @@ export type IdempotencyKeyOptionsSchema = z.infer String(value)); +/** + * Trigger-time named concurrency limits. The charset rule matches what the queue + * concern enforces, so a bad name is a uniform up-front 400 on every path instead + * of a late validation error after the request (or earlier batch items) succeeded. + */ +const TriggerConcurrencyLimitsSchema = z + .string() + .regex(/^[a-zA-Z0-9_-]{1,122}$/, { + message: + "Concurrency limit names are 1-122 characters using only letters, numbers, underscores and hyphens", + }) + .array() + .max(2); + const ExternalDeploymentId = z.preprocess((value) => { // `null` is the opt-out sentinel callers write; treat it as absent, not a validation error. if (value === null) { @@ -343,6 +357,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: TriggerConcurrencyLimitsSchema.optional(), concurrencyKey: ConcurrencyKeySchema.optional(), delay: z.string().or(z.coerce.date()).optional(), idempotencyKey: z @@ -449,6 +473,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: TriggerConcurrencyLimitsSchema.optional(), tags: RunTags.optional(), test: z.boolean().optional(), ttl: z.string().or(z.number().nonnegative().int()).optional(), @@ -1934,6 +1968,12 @@ export const SessionTriggerConfig = z.object({ basePayload: z.record(z.string(), z.unknown()), machine: MachinePresetName.optional(), queue: z.string().max(128).optional(), + /** Named concurrency limits every run holds, replacing the task's declared named limits. + * The charset rule is enforced here so a bad name is rejected before the session row + * persists, instead of surfacing from the trigger after the session already exists. */ + concurrency: TriggerConcurrencyLimitsSchema.optional(), + /** Scopes every run to its own pool under each `perKey` bound it holds. Never defaulted — a session without one shares the keyless pool. */ + concurrencyKey: ConcurrencyKeySchema.optional(), tags: z.array(z.string().max(128)).max(10).optional(), maxAttempts: z.number().int().positive().max(10).optional(), /** Per-run wall-clock cap (seconds). Forwarded to `TaskRunOptions.maxDuration`. */ diff --git a/packages/core/src/v3/schemas/build.ts b/packages/core/src/v3/schemas/build.ts index de9214865b0..a2ce9832a8f 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({ @@ -116,6 +117,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/queues.ts b/packages/core/src/v3/schemas/queues.ts index 78d310174dd..842adabf056 100644 --- a/packages/core/src/v3/schemas/queues.ts +++ b/packages/core/src/v3/schemas/queues.ts @@ -1,4 +1,5 @@ import { z } from "zod/v4"; +import { discriminatedUnion, preprocess } from "../utils/zod.js"; const queueTypes = ["task", "custom"] as const; @@ -13,7 +14,7 @@ export type QueueType = z.infer; export const RetrieveQueueType = z.enum([...queueTypes, "id"]); export type RetrieveQueueType = z.infer; -export const QueueItem = z.object({ +const QueueItemCommon = { /** The queue id, e.g. queue_12345 */ id: z.string(), /** The queue name */ @@ -30,26 +31,60 @@ export const QueueItem = z.object({ queued: z.number(), /** Whether the queue is paused. If it's paused, no new runs will be started. */ paused: z.boolean(), - /** The concurrency limit of the queue */ + /** + * The queue's own concurrency limit. Meaningful on V1 queues only; always + * null on V2 queues (kept on both so existing clients keep parsing). + */ concurrencyLimit: z.number().nullable(), - /** The concurrency limit of the queue */ - concurrency: z - .object({ - /** The effective/current concurrency limit */ - current: z.number().nullable(), - /** The base concurrency limit (default) */ - base: z.number().nullable(), - /** The effective/current concurrency limit */ - override: z.number().nullable(), - /** When the override was applied */ - overriddenAt: z.coerce.date().nullable(), - /** Who overrode the concurrency limit (will be null if overridden via the API) */ - overriddenBy: z.string().nullable(), - }) - .optional(), -}); +}; + +/** + * The queue's `version` discriminates its shape. V1 queues carry their own + * concurrency limit (applied per key when runs pass a `concurrencyKey`, to the + * whole queue when they don't) and its override state. V2 queues are only the + * line runs wait in: concurrency is declared with the task `concurrency` + * option and read or overridden through `concurrencyLimits`. A response from a + * server that predates the discriminator has V1 semantics by definition, so a + * missing `version` defaults to "V1" rather than failing the parse. + */ +const QueueItemUnion = discriminatedUnion("version", [ + z.object({ + ...QueueItemCommon, + version: z.literal("V1"), + /** The queue's concurrency limit override state */ + concurrency: z + .object({ + /** The effective/current concurrency limit */ + current: z.number().nullable(), + /** The base concurrency limit (default) */ + base: z.number().nullable(), + /** The overridden concurrency limit, when an override is active */ + override: z.number().nullable(), + /** When the override was applied */ + overriddenAt: z.coerce.date().nullable(), + /** Who overrode the concurrency limit (will be null if overridden via the API) */ + overriddenBy: z.string().nullable(), + }) + .optional(), + }), + z.object({ + ...QueueItemCommon, + version: z.literal("V2"), + /** Never present on V2 queues; declared so existing `queue.concurrency?.…` + * reads keep compiling across the union and see undefined. */ + concurrency: z.undefined().optional(), + }), +]); + +export const QueueItem = preprocess( + (value) => + value && typeof value === "object" && !("version" in value) + ? { ...value, version: "V1" } + : value, + QueueItemUnion +); -export type QueueItem = z.infer; +export type QueueItem = z.infer; export const ListQueueOptions = z.object({ /** The page number */ @@ -92,3 +127,57 @@ export type QueueTypeName = z.infer; export const RetrieveQueueParam = z.union([z.string(), QueueTypeName]); export type RetrieveQueueParam = z.infer; + +/** One bound of a concurrency limit: the enforced value, the declared base, and + * any active override. `current` is the cap this bound enforces right now (the + * stored value clamped to the environment concurrency limit). Pausing an + * environment or queue stops dequeues without changing its configured bounds, + * so pause state is reported separately, never as a `current` of 0. */ +export const ConcurrencyLimitBound = z.object({ + current: z.number().nullable(), + base: z.number().nullable(), + override: z.number().nullable(), + overriddenAt: z.coerce.date().nullable(), +}); + +export type ConcurrencyLimitBound = z.infer; + +export const ConcurrencyLimitItem = z.object({ + /** The limit's id, starting with `climit_`. */ + id: z.string(), + /** The limit's name, as declared with `concurrencyLimit()` (anonymous inline + * limits use the derived name `task/`). */ + name: z.string(), + /** Caps each concurrencyKey pool; runs without a key share one pool. */ + perKey: ConcurrencyLimitBound, + /** Caps every run holding this limit, keys or not. */ + total: ConcurrencyLimitBound, + /** Runs executing that hold this limit. */ + running: z.number(), + /** Runs that are queued and must clear this limit to execute. */ + queued: z.number(), +}); + +export type ConcurrencyLimitItem = z.infer; + +export const ListConcurrencyLimitOptions = z.object({ + page: z.number().optional(), + perPage: z.number().optional(), +}); + +export type ListConcurrencyLimitOptions = z.infer; + +/** Changes only the given bounds. Zero blocks every run holding the limit, which + * is how a limit is paused; `reset` restores the declared values. */ +export const OverrideConcurrencyLimitRequestBody = z + .object({ + perKey: z.number().int().min(0).max(100000).optional(), + total: z.number().int().min(0).max(100000).optional(), + }) + .refine((body) => body.perKey !== undefined || body.total !== undefined, { + message: "Provide at least one of `perKey` or `total`", + }); + +export type OverrideConcurrencyLimitRequestBody = z.infer< + typeof OverrideConcurrencyLimitRequestBody +>; diff --git a/packages/core/src/v3/schemas/resources.ts b/packages/core/src/v3/schemas/resources.ts index dd7004f4ba8..2ce67bceb27 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/v4"; -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 a87a3fea3d4..6fa0996d046 100644 --- a/packages/core/src/v3/schemas/schemas.ts +++ b/packages/core/src/v3/schemas/schemas.ts @@ -179,10 +179,54 @@ export const QueueManifest = z.object({ * * If this property is omitted, the task can potentially use up the full concurrency of an environment */ concurrencyLimit: z.number().int().min(0).max(100000).optional().nullable(), + /** An optional property that caps the total number of concurrent run executions across ALL + * `concurrencyKey` values of this queue. On a queue with a `concurrencyKey`, `concurrencyLimit` + * applies per key value; this is the ceiling for the whole queue. + * + * Caps keyed and keyless runs together, and requires server-side support. */ + 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(), @@ -198,6 +242,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 9e87f136e2e..fe5e77f94e5 100644 --- a/packages/core/src/v3/types/queues.ts +++ b/packages/core/src/v3/types/queues.ts @@ -31,8 +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 */ + /** + * @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. + */ 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 224b7604db3..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,10 +222,25 @@ type CommonTaskOptions< }); * ``` */ - queue?: { - name?: string; - concurrencyLimit?: 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 @@ -395,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, @@ -811,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/core/src/v3/utils/zod.ts b/packages/core/src/v3/utils/zod.ts index 8f35be3e181..aa74eba269c 100644 --- a/packages/core/src/v3/utils/zod.ts +++ b/packages/core/src/v3/utils/zod.ts @@ -18,3 +18,15 @@ export function discriminatedUnion; } + +/** + * Same declaration-compat trick for `z.preprocess`: newer Zod 4 releases + * declare its result as `ZodPreprocess`, which older `zod/v4` permalinks do + * not export. Erase it to the portable `ZodType` shape in declarations. + */ +export function preprocess( + fn: (value: unknown) => unknown, + schema: T +): z.ZodType, unknown> { + return z.preprocess(fn, schema) as unknown as z.ZodType, unknown>; +} diff --git a/packages/react-hooks/src/hooks/useTaskTrigger.ts b/packages/react-hooks/src/hooks/useTaskTrigger.ts index 979cd87983b..c4f140c5dce 100644 --- a/packages/react-hooks/src/hooks/useTaskTrigger.ts +++ b/packages/react-hooks/src/hooks/useTaskTrigger.ts @@ -73,10 +73,22 @@ export function useTaskTrigger( const payloadPacket = await stringifyIO(payload); + const queueName = options?.queue; + const concurrency = options?.concurrency + ? Array.isArray(options.concurrency) + ? options.concurrency + : [options.concurrency] + : undefined; + + if (concurrency && concurrency.length > 2) { + throw new Error("The concurrency option accepts at most two named limits."); + } + 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/ai.ts b/packages/trigger-sdk/src/v3/ai.ts index a9d40d27091..8113f89ee3b 100644 --- a/packages/trigger-sdk/src/v3/ai.ts +++ b/packages/trigger-sdk/src/v3/ai.ts @@ -94,6 +94,13 @@ import { type TranscriptStorageContext, } from "./transcriptStorage.js"; import { responseAfterCompaction } from "./compactionResponse.js"; +import { ManagedChatResponse, createOrderedChatWriter } from "./managedChatResponse.js"; +import { + convertSteeredMessages, + retainStepMessages, + steeringMarkers, + type SteeringInjection, +} from "./steeringContext.js"; let transcriptStorageOverride: TranscriptStorage | undefined; @@ -156,6 +163,7 @@ import { type SessionTriggerConfigInput, } from "./sessions.js"; import { createTask } from "./shared.js"; +import { triggerConcurrencyBody } from "./concurrency-shared.js"; import { markChatAgentRunForStreamsWarning } from "./streams.js"; import { tracer } from "./tracer.js"; @@ -169,17 +177,24 @@ const METADATA_KEY = "tool.execute.options"; * `ignoreIncompleteToolCalls: true` to prevent failures from * stopped/aborted conversations with partial tool parts. */ -function toModelMessages(messages: UIMessage[]): Promise { +function toModelMessages(messages: UIMessage[], context?: UIMessage[]): Promise { // Pass the resolved per-turn `tools` (if any) so the AI SDK can look up each // tool's `toModelOutput` and re-apply it to prior-turn tool results. Without // `tools` it falls back to JSON-stringifying the raw output (TRI-10149). The // conditional spread keeps the options object byte-identical to the no-tools // path when nothing was declared. const tools = locals.get(chatResolvedToolsKey); - return convertToModelMessages(messages, { - ignoreIncompleteToolCalls: true, - ...(tools ? { tools } : {}), - }); + return convertSteeredMessages( + messages, + async (batch) => + convertToModelMessages(batch, { + ignoreIncompleteToolCalls: true, + ...(tools ? { tools } : {}), + }), + locals.get(chatSteeringInjectionsKey) ?? new Map(), + context ?? locals.get(chatCurrentUIMessagesKey) ?? messages, + context !== undefined || locals.get(chatCurrentUIMessagesKey) !== undefined + ); } export type ToolCallExecutionOptions = { @@ -1307,7 +1322,7 @@ const chatStream: RealtimeDefinedStream = { * `onTurnComplete`'s `responseMessage` and `uiMessages`. * * Non-transient data chunks (`type` starts with `data-`, no `transient: true`) - * are queued for accumulation into the assistant response message. + * are accumulated in emission order into the assistant response message. * Transient or non-data chunks are streamed only (same as `chat.stream`). * * @example @@ -1325,15 +1340,7 @@ const chatResponse = { * response message; everything else is stream-only. */ write(part: UIMessageChunk): void { - queueResponsePart(part); - const { waitUntilComplete } = chatStream.writer({ - spanName: "chat.response.write", - collapsed: true, - execute: ({ write }) => { - write(part); - }, - }); - waitUntilComplete().catch(() => {}); + managedResponse().writeData(part); }, }; @@ -1373,57 +1380,16 @@ export type ChatWriter = { * @internal */ function createLazyChatWriter(): { writer: ChatWriter; flush: () => Promise } { - let writeImpl: ((part: UIMessageChunk) => void) | null = null; - let mergeImpl: ((stream: ReadableStream) => void) | null = null; - let waitPromise: (() => Promise) | null = null; - let resolveExecute: (() => void) | null = null; - let started = false; - const bufferedParts: UIMessageChunk[] = []; - const bufferedStreams: ReadableStream[] = []; - - function ensureInitialized() { - if (started) return; - started = true; - - const executePromise = new Promise((resolve) => { - resolveExecute = resolve; - }); - - const { waitUntilComplete } = chatStream.writer({ - collapsed: true, - spanName: "callback writer", - execute: ({ write, merge }) => { - writeImpl = write; - mergeImpl = merge; - for (const part of bufferedParts.splice(0)) write(part); - for (const stream of bufferedStreams.splice(0)) merge(stream); - return executePromise; - }, - }); - waitPromise = waitUntilComplete; - } - - return { - writer: { - write(part: UIMessageChunk) { - ensureInitialized(); - queueResponsePart(part); - if (writeImpl) writeImpl(part); - else bufferedParts.push(part); - }, - merge(stream: ReadableStream) { - ensureInitialized(); - if (mergeImpl) mergeImpl(stream); - else bufferedStreams.push(stream); - }, - }, - async flush() { - if (resolveExecute) { - resolveExecute(); // Signal execute to complete - await waitPromise!(); // Wait for stream to finish piping - } - }, - }; + return createOrderedChatWriter( + () => (locals.get(chatManagedResponseActiveKey) ? managedResponse() : undefined), + async (stream) => { + const { waitUntilComplete } = chatStream.pipe(stream, { + collapsed: true, + spanName: "callback writer", + }); + await waitUntilComplete(); + } + ); } /** @@ -3702,29 +3668,21 @@ const chatPendingSteerKey = locals.create("chat.pendingSteer"); type PendingSteer = { ui: UIMessage; model: ModelMessage[] }; /** @internal — IDs of messages that were successfully injected via prepareStep */ const chatInjectedMessageIdsKey = locals.create>("chat.injectedMessageIds"); -/** @internal — non-transient data parts queued via chat.response or writer.write() for accumulation into the response message */ -const chatResponsePartsKey = locals.create("chat.responseParts"); - -/** - * Check if a chunk is a non-transient data part that should persist to the response message. - * @internal - */ -function isNonTransientDataPart(part: unknown): boolean { - if (typeof part !== "object" || part === null) return false; - const p = part as Record; - return typeof p.type === "string" && p.type.startsWith("data-") && p.transient !== true; -} - -/** - * Queue a chunk for accumulation into the response message (if it's a non-transient data part). - * Called by `chat.response.write()` and `ChatWriter.write()`. - * @internal - */ -function queueResponsePart(part: unknown): void { - if (!isNonTransientDataPart(part)) return; - const parts = locals.get(chatResponsePartsKey) ?? []; - parts.push(part); - locals.set(chatResponsePartsKey, parts); +const chatManagedResponseKey = locals.create("chat.managedResponse"); +const chatManagedResponseActiveKey = locals.create("chat.managedResponseActive"); +const chatSteeringInjectionsKey = + locals.create>("chat.steeringInjections"); + +function managedResponse(): ManagedChatResponse { + let response = locals.get(chatManagedResponseKey); + if (!response || response.isClosed) { + response = new ManagedChatResponse(async (stream) => { + const { waitUntilComplete } = chatStream.pipe(stream, { spanName: "managed chat response" }); + await waitUntilComplete(); + }); + locals.set(chatManagedResponseKey, response); + } + return response; } /** @@ -4333,6 +4291,7 @@ async function drainSteeringQueue( steps: CompactionStep[], queueOverride?: SteeringQueueEntry[] ): Promise { + if (steps.length === 0) managedResponse().beginGeneration(); const queue = queueOverride ?? locals.get(chatSteeringQueueKey); if (!queue || queue.length === 0) return EMPTY_DRAIN; @@ -4471,30 +4430,24 @@ async function drainSteeringQueue( locals.set(chatPendingSteerKey, pendingSteer); } - // Write injection confirmation chunk to the stream so the frontend - // knows which messages were injected and where in the response. if (injected.length > 0) { - try { - const { waitUntilComplete } = chatStream.writer({ - collapsed: true, - execute: ({ write }) => { - write({ - type: PENDING_MESSAGE_INJECTED_TYPE, - id: generateMessageId(), - data: { - messageIds: claimedUIMessages.map((m) => m.id), - messages: claimedUIMessages.map((m) => ({ - id: m.id, - text: textOfUIMessage(m), - })), - }, - }); - }, - }); - await waitUntilComplete(); - } catch { - /* non-fatal — stream write failed */ - } + const id = generateMessageId(); + const messageIds = claimedUIMessages.map((m) => m.id); + const injections = locals.get(chatSteeringInjectionsKey) ?? new Map(); + injections.set(id, { id, messageIds, messages: injected }); + locals.set(chatSteeringInjectionsKey, injections); + const response = managedResponse(); + // prepareStep can run ahead of the UI consumer. Admit the marker only + // once the actual preceding finish-step chunk has entered this stream. + response.afterStep(steps.length); + response.writeData({ + type: PENDING_MESSAGE_INJECTED_TYPE, + id, + data: { + messageIds, + messages: claimedUIMessages.map((m) => ({ id: m.id, text: textOfUIMessage(m) })), + }, + }); } // Fire onInjected callback @@ -4974,19 +4927,22 @@ function buildManagedStreamTextOptions( ...rest } = options as Record; - const managed = toStreamTextOptions({ - registry, - system: (callerSystem as ToStreamTextOptionsOptions["system"]) ?? agentSystem, - cacheControl, - systemProviderOptions, - /** - * A call site that names `tools` replaces the agent's set rather than - * adding to it, so narrowing the tools for one call still works. Omitting - * `tools` falls back to the agent's, which is what an `onAction` - * regenerate needs: without it a regenerated answer can call nothing. - */ - tools: (tools ?? agentTools) as Record | undefined, - }); + const managed = toStreamTextOptions( + { + registry, + system: (callerSystem as ToStreamTextOptionsOptions["system"]) ?? agentSystem, + cacheControl, + systemProviderOptions, + /** + * A call site that names `tools` replaces the agent's set rather than + * adding to it, so narrowing the tools for one call still works. Omitting + * `tools` falls back to the agent's, which is what an `onAction` + * regenerate needs: without it a regenerated answer can call nothing. + */ + tools: (tools ?? agentTools) as Record | undefined, + }, + false + ); const promptSystem = locals.get(chatPromptKey)?.text; @@ -5022,6 +4978,9 @@ function buildManagedStreamTextOptions( }; } + if (typeof managed.prepareStep === "function") { + managed.prepareStep = retainStepMessages(managed.prepareStep as any); + } return { ...managed, ...rest }; } @@ -5049,7 +5008,10 @@ function createBoundStreamText( return bound as unknown as AiStreamTextFn; } -function toStreamTextOptions(options?: ToStreamTextOptionsOptions): Record { +function toStreamTextOptions( + options?: ToStreamTextOptionsOptions, + retain = true +): Record { const agentDefaults = locals.get(chatAgentManagedConfigKey); if (agentDefaults) { options = { @@ -5288,6 +5250,9 @@ function toStreamTextOptions(options?: ToStreamTextOptionsOptions): Record { */ async function pipeChat( source: UIMessageStreamable | AsyncIterable | ReadableStream, - options?: PipeChatOptions + options?: PipeChatOptions, + capture?: { originalMessages?: UIMessage[] } ): Promise { locals.set(chatPipeCountKey, (locals.get(chatPipeCountKey) ?? 0) + 1); @@ -5536,11 +5502,23 @@ async function pipeChat( // accepts opaque UIMessageStreamable / raw iterables whose element // type we don't know at compile time. Cast — runtime behaviour is // identical (bytes go to session.out either way). - const { waitUntilComplete } = chatStream.pipe( - stream as ReadableStream | AsyncIterable, - pipeOptions - ); - await waitUntilComplete(); + if (capture) { + const response = managedResponse(); + response.seed(capture.originalMessages); + await response.pipe( + stream as ReadableStream | AsyncIterable, + options?.signal + ); + } else { + // Raw/manual pipes intentionally do not feed the managed capture. Never + // leave injection events waiting for finish-step chunks it cannot see. + managedResponse().useRawPipe(); + const { waitUntilComplete } = chatStream.pipe( + stream as ReadableStream | AsyncIterable, + pipeOptions + ); + await waitUntilComplete(); + } } /** @@ -7269,15 +7247,20 @@ function chatAgent< const reconcilePendingSteer = (options?: { /** This turn's model delta, as `onTurnComplete.newMessages` reports it. */ turnNew?: ModelMessage[]; + response?: UIMessage; }): PendingSteer[] => { const pending = locals.get(chatPendingSteerKey); if (!pending || pending.length === 0) return []; locals.set(chatPendingSteerKey, []); - for (const entry of pending) { + const inlineIds = new Set( + steeringMarkers(options?.response ? [options.response] : []).flatMap((m) => m.messageIds) + ); + const standalone = pending.filter((entry) => !inlineIds.has(entry.ui.id)); + for (const entry of standalone) { accumulatedMessages.push(...entry.model); options?.turnNew?.push(...entry.model); } - return pending; + return standalone; }; // Accumulated UI messages for persistence. Mirrors the model accumulator @@ -7395,8 +7378,13 @@ function chatAgent< }); const throughId = opts.messages.at(-1)?.id ?? ""; const queued = locals.get(chatBackgroundQueueKey) ?? []; + const transcriptIds = new Set(opts.messages.map((message) => message.id)); + const markerIds = new Set(steeringMarkers(opts.messages).map((marker) => marker.id)); + const steering = [...(locals.get(chatSteeringInjectionsKey)?.values() ?? [])].filter( + (entry) => markerIds.has(entry.id) || entry.messageIds.some((id) => transcriptIds.has(id)) + ); const runtimeState: TranscriptRuntimeState | null = - laneCompacted || laneInjections.length > 0 || queued.length > 0 + laneCompacted || laneInjections.length > 0 || queued.length > 0 || steering.length > 0 ? { v: 1, ...(laneCompacted @@ -7409,6 +7397,7 @@ function chatAgent< : {}), ...(laneInjections.length > 0 ? { injections: laneInjections } : {}), ...(queued.length > 0 ? { queued: [...queued] } : {}), + ...(steering.length > 0 ? { steering } : {}), } : null; if (runtimeState !== null || persistedStateSet) { @@ -7935,6 +7924,11 @@ function chatAgent< } try { const bootRuntimeState = parseTranscriptRuntimeState(bootTranscriptState); + locals.set( + chatSteeringInjectionsKey, + new Map((bootRuntimeState?.steering ?? []).map((entry) => [entry.id, entry])) + ); + locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); const restored = await restoreModelLane( accumulatedUIMessages, bootRuntimeState, @@ -8479,7 +8473,9 @@ function chatAgent< locals.set(chatCompactionStateKey, undefined); locals.set(chatSteeringQueueKey, []); locals.set(chatPendingBackgroundKey, []); - locals.set(chatResponsePartsKey, []); + await locals.get(chatManagedResponseKey)?.close(); + locals.set(chatManagedResponseKey, undefined); + locals.set(chatManagedResponseActiveKey, true); // NOTE: chatBackgroundQueueKey is NOT reset here — messages injected // by deferred work from the previous turn's onTurnComplete need to // survive into the next turn. The queue is drained before run(). @@ -8663,7 +8659,7 @@ function chatAgent< if (actionOverride) { locals.set(chatOverrideMessagesKey, undefined); accumulatedUIMessages = [...actionOverride] as TUIMessage[]; - accumulatedMessages = await toModelMessages(actionOverride); + accumulatedMessages = await toModelMessages(actionOverride, actionOverride); laneCompacted = false; laneInjections = []; locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); @@ -8848,7 +8844,10 @@ function chatAgent< ) { accumulatedUIMessages.pop(); } - accumulatedMessages = await toModelMessages(accumulatedUIMessages); + accumulatedMessages = await toModelMessages( + accumulatedUIMessages, + accumulatedUIMessages + ); laneCompacted = false; laneInjections = []; } else if (cleanedUIMessages.length > 0) { @@ -8903,7 +8902,10 @@ function chatAgent< logger.warn( "chat.agent: replaced message not found at the model lane tail; reconverting the lane" ); - accumulatedMessages = await toModelMessages(accumulatedUIMessages); + accumulatedMessages = await toModelMessages( + accumulatedUIMessages, + accumulatedUIMessages + ); laneCompacted = false; laneInjections = []; } @@ -9143,7 +9145,10 @@ function chatAgent< if (turnStartOverride) { locals.set(chatOverrideMessagesKey, undefined); accumulatedUIMessages = [...turnStartOverride] as TUIMessage[]; - accumulatedMessages = await toModelMessages(turnStartOverride); + accumulatedMessages = await toModelMessages( + turnStartOverride, + turnStartOverride + ); laneCompacted = false; laneInjections = []; locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); @@ -9232,6 +9237,7 @@ function chatAgent< capturedResponseMessage = lastUI; capturedFinishReason = "stop"; } + managedResponse().seed(accumulatedUIMessages); // Don't call userRun. Don't pipe. Skip directly // to the post-turn flow below. } else { @@ -9304,10 +9310,14 @@ function chatAgent< resolveOnFinish!(); }, }); - await pipeChat(tapUIMessageChunks(uiStream, turnBufferedChunks), { - signal: combinedSignal, - spanName: "stream response", - }); + await pipeChat( + tapUIMessageChunks(uiStream, turnBufferedChunks), + { + signal: combinedSignal, + spanName: "stream response", + }, + { originalMessages: isActionTurn ? undefined : accumulatedUIMessages } + ); } } catch (error) { // Handle AbortError from streamText gracefully @@ -9334,6 +9344,12 @@ function chatAgent< ]); } + capturedResponseMessage = + ((await locals.get(chatManagedResponseKey)?.snapshot()) as + | TUIMessage + | undefined) ?? capturedResponseMessage; + if (capturedResponseMessage) capturedPartialResponse = capturedResponseMessage; + // Capture token usage from the streamText result (if available). // totalUsage is a PromiseLike that resolves after the stream is consumed. // Race with a 2s timeout — on stop-abort the AI SDK's totalUsage @@ -9483,6 +9499,7 @@ function chatAgent< const steerTailThisTurn = reconcilePendingSteer({ turnNew: turnNewModelMessages, + response: capturedResponseMessage, }).reduce((n, e) => n + e.model.length, 0) + reconcilePendingBackground(); // Append the assistant's response (partial or complete) to the accumulator. @@ -9506,15 +9523,6 @@ function chatAgent< id: generateMessageId(), }; } - // Append any non-transient data parts queued via chat.response or writer.write() - const queuedParts = locals.get(chatResponsePartsKey); - if (queuedParts && queuedParts.length > 0) { - capturedResponseMessage = { - ...capturedResponseMessage, - parts: [...capturedResponseMessage.parts, ...queuedParts], - } as TUIMessage; - locals.set(chatResponsePartsKey, []); - } const responseHasContent = capturedResponseMessage.parts.some( (part) => part.type !== "step-start" ); @@ -9595,7 +9603,10 @@ function chatAgent< logger.warn( "chat.agent: replaced response not found at the model lane tail; reconverting the lane" ); - accumulatedMessages = await toModelMessages(accumulatedUIMessages); + accumulatedMessages = await toModelMessages( + accumulatedUIMessages, + accumulatedUIMessages + ); laneCompacted = false; laneInjections = []; } @@ -9613,23 +9624,6 @@ function chatAgent< responseWasSkipped = true; } } - // If there's no captured response (manual pipe mode) but there are - // queued data parts, create a minimal response message to hold them. - if (!capturedResponseMessage) { - const remainingParts = locals.get(chatResponsePartsKey); - if (remainingParts && remainingParts.length > 0) { - capturedResponseMessage = { - id: generateMessageId(), - role: "assistant" as const, - parts: [...remainingParts], - } as TUIMessage; - locals.set(chatResponsePartsKey, []); - accumulatedUIMessages.push(capturedResponseMessage); - turnNewUIMessages.push(capturedResponseMessage); - locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); - } - } - if (capturedResponseMessage) { responseCommitted = true; capturedPartialResponse = capturedResponseMessage; @@ -9816,6 +9810,8 @@ function chatAgent< finishReason: capturedFinishReason, }; + const beforeHookRevision = locals.get(chatManagedResponseKey)?.revision ?? 0; + let beforeHookHistoryEdited = false; // Fire onBeforeTurnComplete — stream is still open so the hook // can write custom chunks to the frontend (e.g. compaction progress). if (onBeforeTurnComplete) { @@ -9829,9 +9825,10 @@ function chatAgent< // Check if the hook replaced messages (compaction or chat.history) const override = locals.get(chatOverrideMessagesKey); if (override) { + beforeHookHistoryEdited = true; locals.set(chatOverrideMessagesKey, undefined); accumulatedUIMessages = [...override] as TUIMessage[]; - accumulatedMessages = await toModelMessages(override); + accumulatedMessages = await toModelMessages(override, override); laneCompacted = false; laneInjections = []; locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); @@ -9851,36 +9848,38 @@ function chatAgent< ); } - // Drain any late response parts added during onBeforeTurnComplete - const lateParts = locals.get(chatResponsePartsKey); - if (lateParts && lateParts.length > 0 && capturedResponseMessage) { + const editedResponse = + beforeHookHistoryEdited && capturedResponseMessage + ? accumulatedUIMessages.find( + (message) => message.id === capturedResponseMessage?.id + ) + : undefined; + const finalManagedResponse = await locals + .get(chatManagedResponseKey) + ?.snapshot( + editedResponse + ? { message: editedResponse, from: beforeHookRevision } + : undefined + ); + if (finalManagedResponse?.parts.some((part) => part.type !== "step-start")) { const idx = accumulatedUIMessages.findIndex( - (m) => m.id === capturedResponseMessage!.id + (m) => m.id === finalManagedResponse.id ); - if (idx !== -1) { - const msg = accumulatedUIMessages[idx]!; - accumulatedUIMessages[idx] = { - ...msg, - parts: [...(msg.parts ?? []), ...lateParts], - } as TUIMessage; - capturedResponseMessage = accumulatedUIMessages[idx] as TUIMessage; - capturedPartialResponse = capturedResponseMessage; - turnCompleteEvent.responseMessage = capturedResponseMessage; - turnCompleteEvent.uiMessages = accumulatedUIMessages; - locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); - } else if (responseWasSkipped) { - capturedResponseMessage = { - ...capturedResponseMessage, - parts: [...(capturedResponseMessage.parts ?? []), ...lateParts], - } as TUIMessage; - accumulatedUIMessages.push(capturedResponseMessage); - turnNewUIMessages.push(capturedResponseMessage); - capturedPartialResponse = capturedResponseMessage; - turnCompleteEvent.responseMessage = capturedResponseMessage; + const finalized = ( + wasStopped ? cleanupAbortedParts(finalManagedResponse) : finalManagedResponse + ) as TUIMessage; + if (idx !== -1 || responseWasSkipped || !capturedResponseMessage) { + if (idx !== -1) accumulatedUIMessages[idx] = finalized; + else accumulatedUIMessages.push(finalized); + const deltaIdx = turnNewUIMessages.findIndex((m) => m.id === finalized.id); + if (deltaIdx !== -1) turnNewUIMessages[deltaIdx] = finalized; + else turnNewUIMessages.push(finalized); + capturedResponseMessage = finalized; + capturedPartialResponse = finalized; + turnCompleteEvent.responseMessage = finalized; turnCompleteEvent.uiMessages = accumulatedUIMessages; locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); } - locals.set(chatResponsePartsKey, []); } settleRecoveredTurn(currentWirePayload); @@ -9906,7 +9905,10 @@ function chatAgent< if (turnCompleteOverride) { locals.set(chatOverrideMessagesKey, undefined); accumulatedUIMessages = [...turnCompleteOverride] as TUIMessage[]; - accumulatedMessages = await toModelMessages(turnCompleteOverride); + accumulatedMessages = await toModelMessages( + turnCompleteOverride, + turnCompleteOverride + ); laneCompacted = false; laneInjections = []; locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages); @@ -10186,6 +10188,7 @@ function chatAgent< : accumulatedUIMessages; let partialResponse: TUIMessage | undefined = + ((await locals.get(chatManagedResponseKey)?.snapshot()) as TUIMessage | undefined) ?? capturedPartialResponse ?? ((await assemblePartialFromChunks(turnBufferedChunks)) as TUIMessage | undefined); if (partialResponse) { @@ -10195,23 +10198,18 @@ function chatAgent< let partialIdx = partialResponse?.id ? erroredUIMessages.findIndex((m) => m.id === partialResponse!.id) : -1; - if (partialResponse && capturedPartialResponse === undefined && partialIdx !== -1) { + if ( + partialResponse && + capturedPartialResponse === undefined && + !locals.get(chatManagedResponseKey)?.continues(partialResponse.id) && + partialIdx !== -1 + ) { partialResponse = undefined; partialIdx = -1; } if (partialResponse && !partialResponse.id) { partialResponse = { ...partialResponse, id: generateMessageId() } as TUIMessage; } - if (partialResponse && !responseCommitted) { - const queuedParts = locals.get(chatResponsePartsKey); - if (queuedParts && queuedParts.length > 0) { - partialResponse = { - ...partialResponse, - parts: [...partialResponse.parts, ...(queuedParts as UIMessage["parts"])], - } as TUIMessage; - locals.set(chatResponsePartsKey, []); - } - } const includePartial = partialResponse != null && !responseCommitted; // What the stream left behind, by content. After `onTurnComplete` the // partial is still unfinished only if the message under its id is @@ -10249,7 +10247,7 @@ function chatAgent< let erroredNewModelMessages: ModelMessage[] = []; - const reconciledSteer = reconcilePendingSteer(); + const reconciledSteer = reconcilePendingSteer({ response: partialResponse }); const backgroundTailThisTurn = reconcilePendingBackground(); if (!responseCommitted) { @@ -10261,17 +10259,9 @@ function chatAgent< * the model received it (what `prepare` produced), matching the * lane. The wire message and partial are converted as before. */ - const steerModelById = new Map( - reconciledSteer.map((e) => [e.ui.id, e.model] as const) + erroredNewModelMessages = await toModelMessages( + erroredNewUIMessages.map(stripProviderMetadata) ); - for (const m of erroredNewUIMessages) { - const recorded = steerModelById.get(m.id); - if (recorded) erroredNewModelMessages.push(...recorded); - else - erroredNewModelMessages.push( - ...(await toModelMessages([stripProviderMetadata(m)])) - ); - } } if (erroredUIMessagesWithPartial !== accumulatedUIMessages) { if (partialIdx === -1) { @@ -10293,7 +10283,10 @@ function chatAgent< logger.warn( "chat.agent: replaced partial not found at the model lane tail; reconverting the lane" ); - accumulatedMessages = await toModelMessages(erroredUIMessagesWithPartial); + accumulatedMessages = await toModelMessages( + erroredUIMessagesWithPartial, + erroredUIMessagesWithPartial + ); laneCompacted = false; laneInjections = []; } @@ -10354,7 +10347,10 @@ function chatAgent< // Convert first: a rejected conversion (a tool's `toModelOutput` // can throw) must leave every lane on the history it had. const overrideUIMessages = [...errorTurnOverride] as TUIMessage[]; - const overrideModelMessages = await toModelMessages(errorTurnOverride); + const overrideModelMessages = await toModelMessages( + errorTurnOverride, + errorTurnOverride + ); erroredUIMessagesWithPartial = overrideUIMessages; accumulatedUIMessages = overrideUIMessages; accumulatedMessages = overrideModelMessages; @@ -10464,6 +10460,11 @@ function chatAgent< // Continue to next iteration of the for loop } finally { turnMsgSub?.off(); + locals.set(chatManagedResponseActiveKey, false); + await locals + .get(chatManagedResponseKey) + ?.close() + .catch(() => {}); } } } finally { @@ -11785,10 +11786,14 @@ async function pipeChatAndCapture( let status: PipeAndCaptureResult["status"] = "complete"; let error: unknown; try { - await pipeChat(tappedStream, { - signal: options?.signal, - spanName: options?.spanName ?? "stream response", - }); + await pipeChat( + tappedStream, + { + signal: options?.signal, + spanName: options?.spanName ?? "stream response", + }, + { originalMessages: options?.originalMessages } + ); // The pipe can drain cleanly on a stop — the source stream just ends // early — so classify by the signal rather than relying on a throw. if (options?.signal?.aborted) { @@ -11815,6 +11820,8 @@ async function pipeChatAndCapture( captured = await assemblePartialFromChunks(bufferedChunks); } + captured = (await locals.get(chatManagedResponseKey)?.snapshot()) ?? captured; + if (!locals.get(chatTurnContextKey)) await locals.get(chatManagedResponseKey)?.close(); return { message: captured, status, @@ -11849,6 +11856,7 @@ class ChatMessageAccumulator { private _handoverRun?: { id: string; run: ModelMessage[] }; private _pendingMessages?: PendingMessagesOptions; private _steeringQueue: SteeringQueueEntry[] = []; + private _pendingSteer: PendingSteer[] = []; constructor(options?: { compaction?: ChatAgentCompactionOptions; @@ -11938,6 +11946,17 @@ class ChatMessageAccumulator { } async addResponse(response: UIMessage): Promise { + const inlineIds = new Set(steeringMarkers([response]).flatMap((m) => m.messageIds)); + for (const entry of this._pendingSteer.splice(0)) { + if (!inlineIds.has(entry.ui.id)) continue; + // absorbSteering is public and updates context immediately. Move just + // those exact appended objects into the response's chronological run; + // never rebuild a possibly compacted lane from the UI transcript. + for (const message of entry.model) { + const index = this.modelMessages.lastIndexOf(message); + if (index !== -1) this.modelMessages.splice(index, 1); + } + } if (!response.id) { response = { ...response, id: generateMessageId() }; } @@ -12015,9 +12034,10 @@ class ChatMessageAccumulator { this.uiMessages.push(...fresh); // Record what the model received. Only when the whole batch is new is // `injected` known to describe exactly these messages. - this.modelMessages.push( - ...(injected && fresh.length === claimed.length ? injected : await toModelMessages(fresh)) - ); + const model = + injected && fresh.length === claimed.length ? injected : await toModelMessages(fresh); + this.modelMessages.push(...model); + for (const ui of fresh) this._pendingSteer.push({ ui, model: modelFormOf(ui, fresh, model) }); } /** @@ -12044,36 +12064,38 @@ class ChatMessageAccumulator { const pm = this._pendingMessages; const queue = this._steeringQueue; - return async ({ messages, steps }) => { - let resultMessages: ModelMessage[] | undefined; + return retainStepMessages( + async ({ messages, steps }: { messages: ModelMessage[]; steps: CompactionStep[] }) => { + let resultMessages: ModelMessage[] | undefined; - // 1. Compaction - if (comp) { - const result = await chatCompact(messages, steps, { - shouldCompact: comp.shouldCompact, - summarize: (msgs) => comp.summarize({ messages: msgs, source: "inner" }), - }); - if (result.type !== "skipped") { - resultMessages = result.messages; + // 1. Compaction + if (comp) { + const result = await chatCompact(messages, steps, { + shouldCompact: comp.shouldCompact, + summarize: (msgs) => comp.summarize({ messages: msgs, source: "inner" }), + }); + if (result.type !== "skipped") { + resultMessages = result.messages; + } } - } - // 2. Pending message injection - if (pm && queue.length > 0) { - const { injected, claimed } = await drainSteeringQueue( - pm, - resultMessages ?? messages, - steps, - queue - ); - await this.absorbSteering(claimed, injected); - if (injected.length > 0) { - resultMessages = [...(resultMessages ?? messages), ...injected]; + // 2. Pending message injection + if (pm) { + const { injected, claimed } = await drainSteeringQueue( + pm, + resultMessages ?? messages, + steps, + queue + ); + await this.absorbSteering(claimed, injected); + if (injected.length > 0) { + resultMessages = [...(resultMessages ?? messages), ...injected]; + } } - } - return resultMessages ? { messages: resultMessages } : undefined; - }; + return resultMessages ? { messages: resultMessages } : undefined; + } + ); } /** @@ -12529,7 +12551,9 @@ function createChatSession( stop.reset(); // Reset per-turn state - locals.set(chatResponsePartsKey, []); + await locals.get(chatManagedResponseKey)?.close(); + locals.set(chatManagedResponseKey, undefined); + locals.set(chatManagedResponseActiveKey, true); // Set up steering queue and pending messages config in locals // so toStreamTextOptions() auto-injects prepareStep for steering const turnSteeringQueue: SteeringQueueEntry[] = []; @@ -12693,11 +12717,6 @@ function createChatSession( if (captured.status === "error") { if (captured.message) { const partial = cleanupAbortedParts(captured.message); - const queuedParts = locals.get(chatResponsePartsKey); - if (queuedParts && queuedParts.length > 0) { - (partial as any).parts = [...(partial.parts ?? []), ...queuedParts]; - locals.set(chatResponsePartsKey, []); - } await accumulator.addResponse(partial); } throw captured.error; @@ -12714,24 +12733,7 @@ function createChatSession( stop.signal.aborted && !runSignal.aborted ? cleanupAbortedParts(response) : response; - // Append any non-transient data parts queued via chat.response or writer.write() - const queuedParts = locals.get(chatResponsePartsKey); - if (queuedParts && queuedParts.length > 0) { - (cleaned as any).parts = [...(cleaned.parts ?? []), ...queuedParts]; - locals.set(chatResponsePartsKey, []); - } await accumulator.addResponse(cleaned); - } else { - // No response (manual pipe mode) but there are queued data parts - const queuedParts = locals.get(chatResponsePartsKey); - if (queuedParts && queuedParts.length > 0) { - await accumulator.addResponse({ - id: generateMessageId(), - role: "assistant" as const, - parts: queuedParts as UIMessage["parts"], - }); - locals.set(chatResponsePartsKey, []); - } } // Capture token usage from the streamText result. Race with a 2s @@ -12829,15 +12831,6 @@ function createChatSession( }, async addResponse(response: UIMessage) { - // Append any non-transient data parts queued via chat.response or writer.write() - const queuedParts = locals.get(chatResponsePartsKey); - if (queuedParts && queuedParts.length > 0) { - response = { - ...response, - parts: [...(response.parts ?? []), ...(queuedParts as UIMessage["parts"])], - }; - locals.set(chatResponsePartsKey, []); - } await accumulator.addResponse(response); }, @@ -12851,41 +12844,43 @@ function createChatSession( const hasPending = !!sessionPendingMessages; if (!hasCompaction && !hasPending) return undefined; - return async ({ - messages: stepMsgs, - steps, - }: { - messages: ModelMessage[]; - steps: CompactionStep[]; - }) => { - let resultMessages: ModelMessage[] | undefined; - - if (sessionCompaction) { - const compactResult = await chatCompact(stepMsgs, steps, { - shouldCompact: sessionCompaction.shouldCompact, - summarize: (msgs) => - sessionCompaction.summarize({ messages: msgs, source: "inner" }), - }); - if (compactResult.type !== "skipped") { - resultMessages = compactResult.messages; + return retainStepMessages( + async ({ + messages: stepMsgs, + steps, + }: { + messages: ModelMessage[]; + steps: CompactionStep[]; + }) => { + let resultMessages: ModelMessage[] | undefined; + + if (sessionCompaction) { + const compactResult = await chatCompact(stepMsgs, steps, { + shouldCompact: sessionCompaction.shouldCompact, + summarize: (msgs) => + sessionCompaction.summarize({ messages: msgs, source: "inner" }), + }); + if (compactResult.type !== "skipped") { + resultMessages = compactResult.messages; + } } - } - if (sessionPendingMessages) { - const { injected, claimed } = await drainSteeringQueue( - sessionPendingMessages, - resultMessages ?? stepMsgs, - steps, - turnSteeringQueue - ); - await accumulator.absorbSteering(claimed, injected); - if (injected.length > 0) { - resultMessages = [...(resultMessages ?? stepMsgs), ...injected]; + if (sessionPendingMessages) { + const { injected, claimed } = await drainSteeringQueue( + sessionPendingMessages, + resultMessages ?? stepMsgs, + steps, + turnSteeringQueue + ); + await accumulator.absorbSteering(claimed, injected); + if (injected.length > 0) { + resultMessages = [...(resultMessages ?? stepMsgs), ...injected]; + } } - } - return resultMessages ? { messages: resultMessages } : undefined; - }; + return resultMessages ? { messages: resultMessages } : undefined; + } + ); }, }; @@ -13330,6 +13325,12 @@ function createChatStartSessionAction( params.clientData !== undefined ? { metadata: params.clientData } : {}; const maxAttempts = params.triggerConfig?.maxAttempts ?? options?.triggerConfig?.maxAttempts; const maxDuration = params.triggerConfig?.maxDuration ?? options?.triggerConfig?.maxDuration; + const concurrency = + params.triggerConfig?.concurrency !== undefined + ? params.triggerConfig.concurrency + : options?.triggerConfig?.concurrency; + const concurrencyKey = + params.triggerConfig?.concurrencyKey ?? options?.triggerConfig?.concurrencyKey; const idleTimeoutInSeconds = params.triggerConfig?.idleTimeoutInSeconds ?? options?.triggerConfig?.idleTimeoutInSeconds; @@ -13355,6 +13356,8 @@ function createChatStartSessionAction( ...(options?.triggerConfig?.queue || params.triggerConfig?.queue ? { queue: params.triggerConfig?.queue ?? options?.triggerConfig?.queue } : {}), + ...(concurrency !== undefined ? triggerConcurrencyBody(concurrency) : {}), + ...(concurrencyKey !== undefined ? { concurrencyKey } : {}), tags, ...(maxAttempts !== undefined ? { maxAttempts } : {}), ...(maxDuration !== undefined ? { maxDuration } : {}), @@ -13788,6 +13791,9 @@ async function writeTurnCompleteChunk( _chatId?: string, publicAccessToken?: string ): Promise { + locals.set(chatManagedResponseActiveKey, false); + const response = locals.get(chatManagedResponseKey); + if (response) await response.close(); const session = getChatSession(); // A handover-prepare boot claims the handover kinds so a signal arriving diff --git a/packages/trigger-sdk/src/v3/chat-client.ts b/packages/trigger-sdk/src/v3/chat-client.ts index d78fb2edcc7..b23748e926f 100644 --- a/packages/trigger-sdk/src/v3/chat-client.ts +++ b/packages/trigger-sdk/src/v3/chat-client.ts @@ -34,6 +34,7 @@ import { slimSubmitMessageForWire, waitBeforeEofResubscribe, } from "./ai-shared.js"; +import { triggerConcurrencyBody } from "./concurrency-shared.js"; import { sessions, type SessionTriggerConfigInput } from "./sessions.js"; // ─── Type inference ──────────────────────────────────────────────── @@ -698,6 +699,12 @@ export class AgentChat { }, ...(this.triggerConfigDefault?.machine ? { machine: this.triggerConfigDefault.machine } : {}), ...(this.triggerConfigDefault?.queue ? { queue: this.triggerConfigDefault.queue } : {}), + ...(this.triggerConfigDefault?.concurrency !== undefined + ? triggerConcurrencyBody(this.triggerConfigDefault.concurrency) + : {}), + ...(this.triggerConfigDefault?.concurrencyKey !== undefined + ? { concurrencyKey: this.triggerConfigDefault.concurrencyKey } + : {}), tags: chatRunTags(this.chatId, this.triggerConfigDefault?.tags), ...(this.triggerConfigDefault?.maxAttempts !== undefined ? { maxAttempts: this.triggerConfigDefault.maxAttempts } diff --git a/packages/trigger-sdk/src/v3/chat-server.ts b/packages/trigger-sdk/src/v3/chat-server.ts index e59075c76df..0981e53906a 100644 --- a/packages/trigger-sdk/src/v3/chat-server.ts +++ b/packages/trigger-sdk/src/v3/chat-server.ts @@ -50,6 +50,7 @@ * `ai.ts` statically imports `agentSkillsRuntime` (which uses `node:` * builtins unfit for some serverless runtimes) and the heavy task * runtime. Allowed imports: `./ai-shared.js`, `./chat-client.js`, + * `./concurrency-shared.js` (dependency-free validation helpers), * `@trigger.dev/core/v3` (api client), `ai` (types + lightweight * helpers like `stepCountIs` / `convertToModelMessages`). */ @@ -79,6 +80,7 @@ import type { } from "ai"; import type { ChatInputChunk, ChatTaskWirePayload } from "./ai-shared.js"; import { chatRunTags } from "./ai-shared.js"; +import { triggerConcurrencyBody } from "./concurrency-shared.js"; import { withResolvedExternalDeploymentId } from "./externalDeploymentId.js"; import type { SessionTriggerConfigInput } from "./sessions.js"; @@ -627,6 +629,12 @@ async function openHandoverSession(opts: { }, ...(opts.triggerConfig?.machine ? { machine: opts.triggerConfig.machine } : {}), ...(opts.triggerConfig?.queue ? { queue: opts.triggerConfig.queue } : {}), + ...(opts.triggerConfig?.concurrency !== undefined + ? triggerConcurrencyBody(opts.triggerConfig.concurrency) + : {}), + ...(opts.triggerConfig?.concurrencyKey !== undefined + ? { concurrencyKey: opts.triggerConfig.concurrencyKey } + : {}), tags, ...(opts.triggerConfig?.maxAttempts !== undefined ? { maxAttempts: opts.triggerConfig.maxAttempts } diff --git a/packages/trigger-sdk/src/v3/compactionResponse.ts b/packages/trigger-sdk/src/v3/compactionResponse.ts index c2a76c186e2..0d1f15e96e8 100644 --- a/packages/trigger-sdk/src/v3/compactionResponse.ts +++ b/packages/trigger-sdk/src/v3/compactionResponse.ts @@ -31,5 +31,9 @@ export function responseAfterCompaction( break; } } + // Injection/data events following the last summarized step precede the + // next step-start. They happened after compaction and belong to the suffix, + // including a stop before that next step ever started. + while (start > 0 && response.parts[start - 1]?.type.startsWith("data-")) start--; return { ...response, parts: response.parts.slice(start) }; } diff --git a/packages/trigger-sdk/src/v3/concurrency-shared.ts b/packages/trigger-sdk/src/v3/concurrency-shared.ts new file mode 100644 index 00000000000..c147cdcee12 --- /dev/null +++ b/packages/trigger-sdk/src/v3/concurrency-shared.ts @@ -0,0 +1,36 @@ +/** + * Concurrency helpers with no runtime dependencies, importable from the lean + * browser and route-handler entrypoints (chat-client, chat-server) without + * pulling the task runtime's module graph into those bundles. + */ + +/** + * 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. + */ +export function triggerConcurrencyBody(concurrency: string | string[] | undefined): { + concurrency?: string[]; +} { + if (concurrency === undefined) { + 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 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.` + ); + } +} diff --git a/packages/trigger-sdk/src/v3/concurrencyLimits.ts b/packages/trigger-sdk/src/v3/concurrencyLimits.ts new file mode 100644 index 00000000000..c67beeb48c1 --- /dev/null +++ b/packages/trigger-sdk/src/v3/concurrencyLimits.ts @@ -0,0 +1,152 @@ +import type { + ApiPromise, + ApiRequestOptions, + ConcurrencyLimitItem, + ListConcurrencyLimitOptions, + OffsetLimitPagePromise, + OverrideConcurrencyLimitRequestBody, +} from "@trigger.dev/core/v3"; +import { + accessoryAttributes, + apiClientManager, + flattenAttributes, + mergeRequestOptions, +} from "@trigger.dev/core/v3"; +import { tracer } from "./tracer.js"; + +/** + * Lists the environment's named concurrency limits (anonymous inline limits appear + * under their derived `task/` names). + * + * @param options - The list options + * @param options.page - The page number + * @param options.perPage - The number of limits per page + * @returns The list of concurrency limits + */ +export function list( + options?: ListConcurrencyLimitOptions, + requestOptions?: ApiRequestOptions +): OffsetLimitPagePromise { + const apiClient = apiClientManager.clientOrThrow(); + + const $requestOptions = mergeRequestOptions( + { + tracer, + name: "concurrencyLimits.list()", + icon: "queue", + }, + requestOptions + ); + + return apiClient.listConcurrencyLimits(options, $requestOptions); +} + +/** + * Retrieves a concurrency limit by name. + * + * @example + * + * ```ts + * const limit = await concurrencyLimits.retrieve("openai"); + * console.log(limit.running, limit.queued, limit.total.current); + * ``` + * @param name - The limit's name, as declared with `concurrencyLimit()` + * @returns The concurrency limit + */ +export function retrieve( + name: string, + requestOptions?: ApiRequestOptions +): ApiPromise { + const apiClient = apiClientManager.clientOrThrow(); + + const $requestOptions = mergeRequestOptions( + { + tracer, + name: "concurrencyLimits.retrieve()", + icon: "queue", + attributes: { + ...flattenAttributes({ name }), + ...accessoryAttributes({ + items: [{ text: name, variant: "normal" }], + style: "codepath", + }), + }, + }, + requestOptions + ); + + return apiClient.retrieveConcurrencyLimit(name, $requestOptions); +} + +/** + * Overrides a concurrency limit's bounds. Only the given fields change; the + * declared values are kept as the base and restored by `reset`. Overriding + * `total` to `0` blocks every run holding the limit, which is how a limit is + * paused. + * + * @example + * + * ```ts + * await concurrencyLimits.override("openai", { total: 50 }); + * await concurrencyLimits.override("openai", { total: 0 }); // pause + * ``` + * @param name - The limit's name + * @param override - The bounds to change (`perKey` and/or `total`) + * @returns The updated concurrency limit + */ +export function override( + name: string, + override: OverrideConcurrencyLimitRequestBody, + requestOptions?: ApiRequestOptions +): ApiPromise { + const apiClient = apiClientManager.clientOrThrow(); + + const $requestOptions = mergeRequestOptions( + { + tracer, + name: "concurrencyLimits.override()", + icon: "queue", + attributes: { + ...flattenAttributes({ name, override }), + ...accessoryAttributes({ + items: [{ text: name, variant: "normal" }], + style: "codepath", + }), + }, + }, + requestOptions + ); + + return apiClient.overrideConcurrencyLimit(name, override, $requestOptions); +} + +/** + * Resets a concurrency limit back to its declared values, clearing any override. + * + * @param name - The limit's name + * @returns The updated concurrency limit + */ +export function reset( + name: string, + requestOptions?: ApiRequestOptions +): ApiPromise { + const apiClient = apiClientManager.clientOrThrow(); + + const $requestOptions = mergeRequestOptions( + { + tracer, + name: "concurrencyLimits.reset()", + icon: "queue", + attributes: { + ...flattenAttributes({ name }), + ...accessoryAttributes({ + items: [{ text: name, variant: "normal" }], + style: "codepath", + }), + }, + }, + requestOptions + ); + + return apiClient.resetConcurrencyLimit(name, $requestOptions); +} diff --git a/packages/trigger-sdk/src/v3/createStartSessionAction.test.ts b/packages/trigger-sdk/src/v3/createStartSessionAction.test.ts index b06f15b2ff1..35881d33a76 100644 --- a/packages/trigger-sdk/src/v3/createStartSessionAction.test.ts +++ b/packages/trigger-sdk/src/v3/createStartSessionAction.test.ts @@ -161,6 +161,96 @@ describe("chat.createStartSessionAction — runtime", () => { expect(lastStartBody?.triggerConfig.ttl).toBe("2m"); }); + it("forwards concurrency and concurrencyKey from triggerConfig, with per-call precedence", async () => { + installStartFixture(); + + const start = chat.createStartSessionAction("fake-chat", { + triggerConfig: { + concurrency: ["chats"], + concurrencyKey: "tenant-default", + }, + }); + await start({ + chatId: "chat-conc", + triggerConfig: { concurrencyKey: "tenant-42" }, + }); + + expect(lastStartBody?.triggerConfig.concurrency).toEqual(["chats"]); + expect(lastStartBody?.triggerConfig.concurrencyKey).toBe("tenant-42"); + }); + + it("per-call concurrency wins over the action default, and an empty array clears it", async () => { + installStartFixture(); + + const start = chat.createStartSessionAction("fake-chat", { + triggerConfig: { concurrency: ["chats"] }, + }); + + await start({ + chatId: "chat-conc-override", + triggerConfig: { concurrency: ["priority"] }, + }); + expect(lastStartBody?.triggerConfig.concurrency).toEqual(["priority"]); + + await start({ + chatId: "chat-conc-clear", + triggerConfig: { concurrency: [] }, + }); + expect(lastStartBody?.triggerConfig.concurrency).toEqual([]); + }); + + it("never defaults concurrencyKey from the chatId", async () => { + installStartFixture(); + + const start = chat.createStartSessionAction("fake-chat"); + await start({ chatId: "chat-no-key" }); + + expect(lastStartBody?.triggerConfig.concurrencyKey).toBeUndefined(); + expect(lastStartBody?.triggerConfig.concurrency).toBeUndefined(); + }); + + it("rejects invalid trigger-time limit names before any network call", async () => { + installStartFixture(); + + const start = chat.createStartSessionAction("fake-chat", { + triggerConfig: { concurrency: ["not a valid name!"] }, + }); + + await expect(start({ chatId: "chat-bad-limit" })).rejects.toThrow( + /letters, numbers, underscores and hyphens/ + ); + expect(lastStartBody).toBeUndefined(); + }); + + it("rejects an empty-string concurrency instead of silently dropping it", async () => { + installStartFixture(); + + const emptyDefault = chat.createStartSessionAction("fake-chat", { + triggerConfig: { concurrency: "" as unknown as string[] }, + }); + await expect(emptyDefault({ chatId: "chat-empty-limit" })).rejects.toThrow(/non-empty strings/); + + const emptyPerCall = chat.createStartSessionAction("fake-chat"); + await expect( + emptyPerCall({ + chatId: "chat-empty-limit-2", + triggerConfig: { concurrency: "" as unknown as string[] }, + }) + ).rejects.toThrow(/non-empty strings/); + + const nullPerCall = chat.createStartSessionAction("fake-chat", { + triggerConfig: { concurrency: ["chats"] }, + }); + await expect( + nullPerCall({ + chatId: "chat-null-limit", + triggerConfig: { concurrency: null as unknown as string[] }, + }) + ).rejects.toThrow(/non-empty strings/); + + expect(lastStartBody).toBeUndefined(); + }); + it("server-mints override tokens for additional API keys", async () => { const requests: Array<{ url: string; body: unknown }> = []; const start = chat.createStartSessionAction("fake-chat", { diff --git a/packages/trigger-sdk/src/v3/index.ts b/packages/trigger-sdk/src/v3/index.ts index 4bdb582d7ae..64c6f85735b 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"; @@ -65,6 +65,7 @@ export { } from "./deployments.js"; export * as envvars from "./envvars.js"; export * as queues from "./queues.js"; +export * as concurrencyLimits from "./concurrencyLimits.js"; export type { ImportEnvironmentVariablesParams } from "./envvars.js"; export { configure, auth } from "./auth.js"; diff --git a/packages/trigger-sdk/src/v3/managedChatResponse.test.ts b/packages/trigger-sdk/src/v3/managedChatResponse.test.ts new file mode 100644 index 00000000000..eb735bc2004 --- /dev/null +++ b/packages/trigger-sdk/src/v3/managedChatResponse.test.ts @@ -0,0 +1,545 @@ +import { + convertToModelMessages, + readUIMessageStream, + type UIMessage, + type UIMessageChunk, + type ModelMessage, +} from "ai"; +import { describe, it, expect } from "vitest"; +import { ManagedChatResponse, createOrderedChatWriter } from "./managedChatResponse.js"; +import { + convertSteeredMessages, + retainStepMessages, + type SteeringInjection, +} from "./steeringContext.js"; +import { + normalizeRuntimeStateForWindow, + parseTranscriptRuntimeState, + restoreModelLane, +} from "./transcriptStorage.js"; +import { responseAfterCompaction } from "./compactionResponse.js"; + +const user = (id: string, text: string): UIMessage => ({ + id, + role: "user", + parts: [{ type: "text", text }], +}); +const modelUser = (content: string): ModelMessage => ({ role: "user", content }); +const marker: UIMessageChunk = { + type: "data-pending-message-injected", + id: "injection-1", + data: { messageIds: ["steer"], messages: [{ id: "steer", text: "Only Platform" }] }, +}; +const firstStep: UIMessageChunk[] = [ + { type: "start", messageId: "answer" }, + { type: "start-step" }, + { type: "tool-input-available", toolCallId: "lookup-1", toolName: "lookup", input: {} }, + { + type: "tool-output-available", + toolCallId: "lookup-1", + output: { projects: ["Platform", "Website"] }, + }, + { type: "finish-step" }, +]; +const lastStep: UIMessageChunk[] = [ + { type: "start-step" }, + { type: "text-start", id: "text" }, + { type: "text-delta", id: "text", delta: "Platform is ready." }, + { type: "text-end", id: "text" }, + { type: "finish-step" }, + { type: "finish" }, +]; +function source(chunks: UIMessageChunk[]) { + return new ReadableStream({ + start(c) { + for (const chunk of chunks) c.enqueue(chunk); + c.close(); + }, + }); +} +function response() { + const emitted: UIMessageChunk[] = []; + const capture = new ManagedChatResponse(async (stream) => { + for await (const chunk of stream) emitted.push(chunk); + }); + return { capture, emitted }; +} +async function reduce(chunks: UIMessageChunk[], message?: UIMessage) { + let result: UIMessage | undefined; + for await (const update of readUIMessageStream({ stream: source(chunks), message })) + result = update; + return result; +} +const convert = async (messages: UIMessage[]) => + convertToModelMessages(messages, { ignoreIncompleteToolCalls: true }); + +describe("ordered managed response", () => { + it("matches the browser and persistence, even when prepareStep outruns the UI consumer", async () => { + const { capture, emitted } = response(); + capture.afterStep(1); + capture.writeData(marker); + capture.writeData({ type: "data-position", id: "custom", data: "between steps" }); + await capture.pipe(source([...firstStep, ...lastStep])); + const saved = await capture.snapshot(); + await capture.close(); + expect(await reduce(emitted)).toEqual(saved); + expect(saved?.parts.map((p) => p.type)).toEqual([ + "step-start", + "tool-lookup", + "data-pending-message-injected", + "data-position", + "step-start", + "text", + ]); + expect(emitted.indexOf(marker)).toBe(emitted.findIndex((c) => c.type === "finish-step") + 1); + }); + + it.each(["end", "error", "abort", "close"])( + "preserves boundary-gated data when the source terminates early: %s", + async (termination) => { + const { capture, emitted } = response(); + const custom: UIMessageChunk = { type: "data-position", data: "accepted" }; + capture.afterStep(1); + capture.writeData(marker); + capture.writeData(custom); + if (termination === "close") { + await capture.close(); + } else { + const abort = new AbortController(); + let next = 0; + const early = new ReadableStream({ + pull(controller) { + if (next < firstStep.length - 1) controller.enqueue(firstStep[next++]!); + else if (termination === "error") controller.error(new Error("early failure")); + else if (termination === "abort") abort.abort(new Error("stopped")); + else controller.close(); + }, + }); + const work = capture.pipe(early, abort.signal); + if (termination === "error") await expect(work).rejects.toThrow("early failure"); + else if (termination === "abort") await work.catch(() => {}); + else await work; + } + // SDK callers snapshot before closing the response on source termination. + const saved = await capture.snapshot(); + await capture.close(); + expect(saved?.parts.slice(-2)).toEqual([ + { type: marker.type, id: "injection-1", data: marker.data }, + { type: custom.type, data: "accepted" }, + ]); + expect(emitted.filter((chunk) => chunk === marker || chunk === custom)).toEqual([ + marker, + custom, + ]); + expect(await reduce(emitted)).toEqual(saved); + } + ); + + it("aborts a pending async iterator without waiting for its cleanup", async () => { + const { capture } = response(); + const abort = new AbortController(); + let started!: () => void; + const waiting = new Promise((resolve) => { + started = resolve; + }); + let returned = false; + const blocked: AsyncIterable = { + [Symbol.asyncIterator]() { + return { + next() { + started(); + return new Promise>(() => {}); + }, + return() { + returned = true; + return new Promise>(() => {}); + }, + }; + }, + }; + const work = capture.pipe(blocked, abort.signal); + await waiting; + const reason = new Error("stop requested"); + abort.abort(reason); + let timer: ReturnType | undefined; + try { + const outcome = await Promise.race([ + work.then( + () => "completed", + (error) => error + ), + new Promise((resolve) => { + timer = setTimeout(() => resolve("still blocked"), 200); + }), + ]); + expect(outcome).toBe(reason); + expect(returned).toBe(true); + } finally { + clearTimeout(timer); + await capture.close(); + } + }); + + it.each([true, false])("preserves mixed hook-stream order (managed turn: %s)", async (active) => { + const emitted: UIMessageChunk[] = []; + let publishers = 0; + const publish = async (stream: ReadableStream) => { + publishers++; + // Different publishers would be able to overtake this asynchronous reader. + await new Promise((resolve) => setTimeout(resolve, 5)); + for await (const chunk of stream) emitted.push(chunk); + }; + const capture = new ManagedChatResponse(publish); + const { writer, flush } = createOrderedChatWriter( + () => (active ? capture : undefined), + publish + ); + const mixed: UIMessageChunk[] = [ + { type: "data-status", id: "status", data: "loading" }, + { type: "start-step" }, + { type: "data-note", data: "before text" }, + { type: "text-start", id: "hook-text" }, + { type: "text-delta", id: "hook-text", delta: "stream only" }, + { type: "data-status", id: "status", data: "done" }, + { type: "text-end", id: "hook-text" }, + { type: "finish-step" }, + { type: "data-progress", transient: true, data: "transient" }, + ]; + if (active) capture.afterStep(1); + writer.merge(source(mixed)); + await flush(); + if (active) await capture.pipe(source(firstStep)); + await capture.close(); + expect(publishers).toBe(1); + expect(emitted).toEqual(active ? [...firstStep, ...mixed] : mixed); + const saved = await capture.snapshot(); + if (active) { + expect(saved?.parts.filter((part) => part.type.startsWith("data-"))).toEqual([ + { type: "data-status", id: "status", data: "done" }, + { type: "data-note", data: "before text" }, + ]); + expect(saved?.parts.some((part) => part.type === "text")).toBe(false); + } else expect(saved).toBeUndefined(); + }); + + it("orders injection against the current generation after an earlier generation finished", async () => { + const { capture, emitted } = response(); + await capture.pipe(source([...firstStep, ...lastStep])); + capture.beginGeneration(); + capture.afterStep(1); + capture.writeData(marker); + const second = firstStep.slice(1); + await capture.pipe(source(second)); + await capture.close(); + expect(emitted.indexOf(marker)).toBe(emitted.map((c) => c.type).lastIndexOf("finish-step") + 1); + }); + + it("keeps raw-pipe injection events observable without capturing raw model output", async () => { + const { capture, emitted } = response(); + capture.afterStep(1); + capture.writeData(marker); + capture.useRawPipe(); + capture.afterStep(2); + const custom: UIMessageChunk = { type: "data-position", data: "raw owner" }; + capture.writeData(custom); + await capture.close(); + expect(emitted.slice(0, 2)).toEqual([marker, custom]); + expect((await capture.snapshot())?.parts.map((p) => p.type)).toEqual([ + "data-pending-message-injected", + "data-position", + ]); + }); + + it("gives a data-only response the same stable ID live and in storage", async () => { + const { capture, emitted } = response(); + capture.writeData({ type: "data-card", data: "No model needed" }); + const saved = await capture.snapshot(); + await capture.close(); + expect(saved?.id).toBeTruthy(); + expect(await reduce(emitted)).toEqual(saved); + }); + + it("applies late writes to an explicitly edited response without undoing the edit", async () => { + const { capture } = response(); + await capture.pipe(source([...firstStep, ...lastStep])); + const saved = (await capture.snapshot())!; + const revision = capture.revision; + const edited: UIMessage = { + ...saved, + parts: [{ type: "text", text: "Application edited this" }], + }; + capture.writeData({ type: "data-note", data: "late" }); + expect((await capture.snapshot({ message: edited, from: revision }))?.parts).toEqual([ + { type: "text", text: "Application edited this" }, + { type: "data-note", data: "late" }, + ]); + await capture.close(); + }); + + it("updates existing IDs in place across early and late hook writes and excludes transient data", async () => { + const { capture, emitted } = response(); + capture.writeData({ type: "data-status", id: "status", data: "loading" }); + await capture.pipe(source([...firstStep, ...lastStep])); + capture.writeData({ type: "data-status", id: "status", data: "done" }); + capture.writeData({ type: "data-progress", transient: true, data: "ephemeral" }); + const saved = await capture.snapshot(); + await capture.close(); + expect(saved).toEqual(await reduce(emitted)); + expect(saved?.parts[0]).toMatchObject({ type: "data-status", data: "done" }); + expect(saved?.parts.filter((p) => p.type === "data-status")).toHaveLength(1); + expect(JSON.stringify(saved)).not.toContain("ephemeral"); + }); + + it("seeds same-ID continuations without mutating the original assistant", async () => { + const original = (await reduce(firstStep))!; + const before = structuredClone(original); + const { capture, emitted } = response(); + capture.seed([original]); + capture.writeData(marker); + await capture.pipe(source([{ type: "start", messageId: original.id }, ...lastStep])); + const saved = await capture.snapshot(); + await capture.close(); + expect(original).toEqual(before); + expect(saved).toEqual(await reduce(emitted, structuredClone(original))); + expect(saved?.parts.map((p) => p.type)).toEqual([ + "step-start", + "tool-lookup", + "data-pending-message-injected", + "step-start", + "text", + ]); + }); + + it("recovers already emitted data and text when the source errors", async () => { + const { capture, emitted } = response(); + await capture.pipe(source(firstStep)); + capture.writeData(marker); + const chunks = lastStep.slice(0, 3); + const broken = new ReadableStream({ + pull(c) { + const chunk = chunks.shift(); + if (chunk) c.enqueue(chunk); + else c.error(new Error("source failed")); + }, + }); + await expect(capture.pipe(broken)).rejects.toThrow("source failed"); + const saved = await capture.snapshot(); + await capture.close(); + expect(saved).toEqual(await reduce(emitted)); + expect(saved?.parts.at(-1)).toMatchObject({ type: "text", text: "Platform is ready." }); + }); +}); + +describe("steering model context", () => { + const prepared: SteeringInjection = { + id: "injection-1", + messageIds: ["steer"], + messages: [{ role: "system", content: "PRIVATE transformed steering" }], + }; + const forms = new Map([[prepared.id, prepared]]); + + it("preserves overrides over later steps on both resetting and retaining AI SDK versions", async () => { + for (const sdkRetains of [false, true]) { + const first = modelUser("original"); + const steer = modelUser("steer"); + const outputs: ModelMessage[] = [ + { role: "assistant", content: "step one" }, + { role: "assistant", content: "step two" }, + ]; + const prepare = retainStepMessages(async ({ messages, steps }) => + steps.length === 1 ? { messages: [...messages, steer] } : undefined + ); + expect(await prepare({ messages: [first], steps: [] })).toBeUndefined(); + const second = await prepare({ + messages: [first, outputs[0]!], + steps: [{ response: { messages: outputs.slice(0, 1) } }], + }); + expect(second?.messages).toEqual([first, outputs[0], steer]); + const third = await prepare({ + messages: sdkRetains ? [...second!.messages!, outputs[1]!] : [first, ...outputs], + steps: outputs.map((_, i) => ({ response: { messages: outputs.slice(0, i + 1) } })), + }); + expect(third?.messages).toEqual([first, outputs[0], steer, outputs[1]]); + expect( + await prepare({ messages: [modelUser("fresh generation")], steps: [] }) + ).toBeUndefined(); + } + }); + + it("retains compaction instead of reintroducing old response segments", async () => { + const summary = modelUser("summary"); + const old = modelUser("old"); + const next = modelUser("new output"); + const prepare = retainStepMessages(async ({ steps }) => + steps.length === 1 ? { messages: [summary] } : undefined + ); + await prepare({ messages: [old], steps: [{ response: { messages: [old] } }] }); + expect( + ( + await prepare({ + messages: [old, next], + steps: [{ response: { messages: [old] } }, { response: { messages: [old, next] } }], + }) + )?.messages + ).toEqual([summary, next]); + }); + + it("keeps complete tool pairs before transformed steering and restores that order in a fresh worker", async () => { + const answer = (await reduce([...firstStep, marker, ...lastStep]))!; + const ui = [user("original", "all projects"), user("steer", "Only Platform"), answer]; + const state = parseTranscriptRuntimeState( + JSON.parse(JSON.stringify({ v: 1, steering: [prepared] })) + )!; + const restoredForms = new Map(state.steering!.map((entry) => [entry.id, entry])); + const projection = (messages: UIMessage[]) => + convertSteeredMessages(messages, convert, restoredForms, ui); + const model = await projection(ui); + expect(model.map((m) => m.role)).toEqual(["user", "assistant", "tool", "system", "assistant"]); + expect(model[3]).toEqual(prepared.messages[0]); + expect(JSON.stringify(model)).not.toContain("Only Platform"); + expect((await restoreModelLane(ui, state, projection)).messages).toEqual(model); + expect(JSON.stringify(answer)).not.toContain("PRIVATE"); + }); + + it.each([false, true])( + "keeps tools, steering and answers across two later turns (reload runtime state: %s)", + async (restart) => { + const secondTool: UIMessageChunk[] = firstStep + .slice(1) + .map((chunk) => ("toolCallId" in chunk ? { ...chunk, toolCallId: "lookup-2" } : chunk)); + const firstAnswer = (await reduce([...firstStep, marker, ...secondTool, ...lastStep]))!; + let history = [user("original", "all projects"), user("steer", "Only Platform"), firstAnswer]; + const originalAnswer = structuredClone(firstAnswer); + let runtime = { v: 1 as const, steering: [prepared] }; + let model = await convertSteeredMessages(history, convert, forms, history, true); + expect(model.map((m) => m.role)).toEqual([ + "user", + "assistant", + "tool", + "system", + "assistant", + "tool", + "assistant", + ]); + const initialModel = structuredClone(model); + for (let turn = 1; turn <= 2; turn++) { + const incoming = user(`followup-${turn}`, `Follow-up ${turn}: keep the same scope`); + history.push(incoming); + const expectedInput = [...model, ...(await convert([incoming]))]; + if (restart) { + // Exercise the actual persisted state parser and model-lane restore, + // with fresh objects so in-memory references cannot hide state loss. + const saved = JSON.parse(JSON.stringify({ messages: history, state: runtime })); + history = saved.messages; + runtime = parseTranscriptRuntimeState(saved.state) as typeof runtime; + } + const currentForms = new Map(runtime.steering.map((entry) => [entry.id, entry])); + const project = (messages: UIMessage[]) => + convertSteeredMessages(messages, convert, currentForms, history, true); + const restored = await restoreModelLane(history, runtime, project); + expect(restored.messages).toEqual(expectedInput); + expect(restored.messages.slice(0, initialModel.length)).toEqual(initialModel); + const followupTool: UIMessageChunk[] = firstStep.map((chunk) => { + if (chunk.type === "start") return { ...chunk, messageId: `answer-${turn}` }; + return "toolCallId" in chunk + ? { ...chunk, toolCallId: `followup-lookup-${turn}` } + : chunk; + }); + const answer = (await reduce([...followupTool, ...lastStep]))!; + history.push(answer); + model = [...expectedInput, ...(await convert([answer]))]; + expect(await project(history)).toEqual(model); + expect(history.find((message) => message.id === originalAnswer.id)).toEqual(originalAnswer); + expect(new Set(history.map((message) => message.id)).size).toBe(history.length); + expect(model.filter((message) => message.role === "system")).toEqual(prepared.messages); + expect(JSON.stringify(history)).not.toContain("PRIVATE"); + } + } + ); + + it("can reconstruct ordinary steering from UI alone without duplicating its user message", async () => { + const answer = (await reduce([...firstStep, marker, ...lastStep]))!; + const model = await convertSteeredMessages( + [user("steer", "Only Platform"), answer], + convert, + new Map() + ); + expect(model.map((m) => m.role)).toEqual(["assistant", "tool", "user", "assistant"]); + expect(JSON.stringify(model).match(/Only Platform/g)).toHaveLength(1); + }); + + it("does not lose an injection immediately after inner compaction", async () => { + const answer = (await reduce([...firstStep, marker, ...lastStep]))!; + const suffix = responseAfterCompaction(answer, 1); + const model = await convertSteeredMessages([suffix], convert, forms); + expect(model.map((m) => m.role)).toEqual(["system", "assistant"]); + expect(JSON.stringify(model)).not.toContain("Website"); + }); + + it("retains private steering across a trimmed input and subsequent worker restart", async () => { + const answer = (await reduce([...firstStep, marker, ...lastStep]))!; + const history = [answer]; + const runtime = normalizeRuntimeStateForWindow({ v: 1, steering: [prepared] }, history, [ + "steer", + answer.id, + ]); + const restored = parseTranscriptRuntimeState(JSON.parse(JSON.stringify(runtime)))!; + expect(restored.steering).toHaveLength(1); + const retainedForms = new Map(restored.steering!.map((entry) => [entry.id, entry])); + const project = (messages: UIMessage[]) => + convertSteeredMessages(messages, convert, retainedForms, history, true); + const expected = await convertSteeredMessages([answer], convert, forms); + expect((await restoreModelLane(history, restored, project)).messages).toEqual(expected); + const resaved = parseTranscriptRuntimeState( + JSON.parse( + JSON.stringify( + normalizeRuntimeStateForWindow( + restored, + history, + history.map((m) => m.id) + ) + ) + ) + )!; + expect(resaved.steering).toEqual(restored.steering); + history.push(user("next", "Follow-up")); + expect((await restoreModelLane(history, restored, project)).messages).toEqual([ + ...expected, + ...(await convert([history[1]!])), + ]); + expect(JSON.stringify(history)).not.toContain("PRIVATE"); + const pruned = parseTranscriptRuntimeState( + normalizeRuntimeStateForWindow( + restored, + [history[1]!], + history.map((m) => m.id) + ) + )!; + expect(pruned.steering).toEqual([]); + }); + + it("honors an explicit deletion even when the saved marker retains fallback text", async () => { + const answer = (await reduce([...firstStep, marker, ...lastStep]))!; + const history = [user("original", "question"), answer]; + const state = parseTranscriptRuntimeState( + JSON.parse( + JSON.stringify( + normalizeRuntimeStateForWindow( + { v: 1, steering: [prepared] }, + history, + history.map((m) => m.id) + ) + ) + ) + )!; + const savedForms = new Map(state.steering!.map((entry) => [entry.id, entry])); + const model = await convertSteeredMessages(history, convert, savedForms, history, true); + expect(JSON.stringify(model)).not.toContain("PRIVATE"); + expect(JSON.stringify(model)).not.toContain("Only Platform"); + }); + + it("retains a claimed instruction if stopping leaves no captured injection marker", async () => { + const model = await convertSteeredMessages([user("steer", "Only Platform")], convert, forms); + expect(model).toEqual(prepared.messages); + }); +}); diff --git a/packages/trigger-sdk/src/v3/managedChatResponse.ts b/packages/trigger-sdk/src/v3/managedChatResponse.ts new file mode 100644 index 00000000000..e67393da6c2 --- /dev/null +++ b/packages/trigger-sdk/src/v3/managedChatResponse.ts @@ -0,0 +1,231 @@ +import type { UIMessage, UIMessageChunk } from "ai"; +import { generateId, readUIMessageStream } from "../imports/ai-runtime.js"; + +/** One ordered output channel for a turn's managed model and data chunks. */ +export class ManagedChatResponse { + private controller?: ReadableStreamDefaultController; + private completion?: Promise; + private failure?: unknown; + private closed = false; + private chunks: UIMessageChunk[] = []; + private original?: UIMessage; + private fallbackId = generateId(); + private completedSteps = 0; + private dataBoundary = 0; + private generationOffset = 0; + private orderedSteps = true; + private pendingData: Array<{ chunk: UIMessageChunk; capture: boolean; after: number }> = []; + + constructor( + private readonly publish: (stream: ReadableStream) => Promise + ) {} + + seed(messages: readonly UIMessage[] | undefined): void { + // Never mutate the caller's history while processing a continuation. + const last = messages?.at(-1); + this.original = last?.role === "assistant" ? structuredClone(last) : undefined; + } + + get isClosed(): boolean { + return this.closed; + } + get revision(): number { + return this.chunks.length; + } + + continues(id: string): boolean { + return this.original?.id === id; + } + + beginGeneration(): void { + this.generationOffset = this.completedSteps; + this.dataBoundary = this.completedSteps; + } + + afterStep(completedSteps: number): void { + this.dataBoundary = this.orderedSteps ? this.generationOffset + completedSteps : 0; + } + + /** Raw pipes own their output ordering and do not feed this capture. */ + useRawPipe(): void { + this.orderedSteps = false; + this.dataBoundary = 0; + this.flushPendingData(); + } + + private flushPendingData(): void { + this.dataBoundary = this.completedSteps; + for (const next of this.pendingData.splice(0)) this.write(next.chunk, next.capture); + } + + writeData(chunk: UIMessageChunk): void { + const capture = chunk.type.startsWith("data-") && !("transient" in chunk && chunk.transient); + if (this.closed) throw new Error("The managed chat response is closed"); + if (this.dataBoundary > this.completedSteps) { + this.pendingData.push({ chunk, capture, after: this.dataBoundary }); + } else this.write(chunk, capture); + } + + write(chunk: UIMessageChunk, capture = true): void { + if (this.closed) throw new Error("The managed chat response is closed"); + if (this.failure) throw this.failure; + if (!this.controller) { + const stream = new ReadableStream({ + start: (controller) => { + this.controller = controller; + }, + }); + this.completion = this.publish(stream); + // Transport errors are observed by pipe/close; never leave an unhandled + // rejection when a fire-and-forget custom write opened the response. + this.completion.catch((error) => { + this.failure = error; + }); + } + if (capture) this.chunks.push(chunk); + this.controller!.enqueue(chunk); + if (capture && chunk.type === "finish-step") { + this.completedSteps++; + while (this.pendingData[0] && this.pendingData[0].after <= this.completedSteps) { + const next = this.pendingData.shift()!; + this.write(next.chunk, next.capture); + } + } + } + + async pipe( + source: AsyncIterable | ReadableStream, + signal?: AbortSignal + ): Promise { + const reader = source instanceof ReadableStream ? source.getReader() : undefined; + const iterator = reader + ? undefined + : (source as AsyncIterable)[Symbol.asyncIterator](); + let rejectAbort!: (reason: unknown) => void; + const aborted = new Promise((_, reject) => { + rejectAbort = reject; + }); + // A pre-aborted signal or a synchronously throwing iterator can leave this + // promise outside the read race; observe its rejection in either case. + aborted.catch(() => {}); + const cancel = () => rejectAbort(signal?.reason); + signal?.addEventListener("abort", cancel, { once: true }); + try { + while (true) { + signal?.throwIfAborted(); + if (this.failure) throw this.failure; + const next = reader ? reader.read() : iterator!.next(); + const result = signal ? await Promise.race([next, aborted]) : await next; + signal?.throwIfAborted(); + if (result.done) break; + this.write(result.value); + } + } finally { + signal?.removeEventListener("abort", cancel); + try { + if (signal?.aborted) { + // Async generators can queue return() behind a blocked next(). Stop + // the turn without waiting for provider cleanup, and observe failures. + try { + void Promise.resolve( + reader ? reader.cancel(signal.reason) : iterator?.return?.() + ).catch(() => {}); + } catch { + // A synchronous cleanup failure must not replace the abort reason. + } + } else if (iterator) await iterator.return?.(); + } finally { + reader?.releaseLock(); + // No further finish-step can arrive from this source. Preserve accepted + // data at the end of the available partial before callers snapshot it. + this.flushPendingData(); + } + } + } + + /** Native reduction of the exact admitted chunks, including same-ID updates. */ + async snapshot(options?: { message: UIMessage; from: number }): Promise { + if (!this.chunks.length) return options?.message; + const chunks = this.chunks.slice(options?.from ?? 0); + let message: UIMessage | undefined = options?.message; + const original = options?.message ?? this.original; + const stream = new ReadableStream({ + start(controller) { + for (const chunk of chunks) controller.enqueue(chunk); + controller.close(); + }, + }); + for await (const update of readUIMessageStream({ + stream, + ...(original ? { message: structuredClone(original) } : {}), + // An error chunk must leave previously streamed content recoverable. + terminateOnError: false, + })) + message = update; + return message ? { ...message, id: message.id || this.fallbackId } : undefined; + } + + async close(): Promise { + if (!this.closed) { + this.flushPendingData(); + // A data-only response still needs the same server ID on both sides. + if ( + !this.original && + this.chunks.length && + !this.chunks.some((chunk) => chunk.type === "start") + ) { + this.write({ type: "start", messageId: this.fallbackId }); + } + this.closed = true; + this.controller?.close(); + } + await this.completion; + } +} + +/** Hook streams share a publisher; only persistent data joins managed capture. */ +export function createOrderedChatWriter( + activeResponse: () => ManagedChatResponse | undefined, + publish: (stream: ReadableStream) => Promise +) { + const raw = new ManagedChatResponse(publish); + const merging: Promise[] = []; + const writer = { + write(part: UIMessageChunk) { + const managed = activeResponse(); + if (managed) managed.writeData(part); + else raw.write(part, false); + }, + merge(stream: ReadableStream) { + const work = (async () => { + const reader = stream.getReader(); + try { + while (true) { + const { done, value } = await reader.read(); + if (done) break; + writer.write(value); + } + } finally { + reader.releaseLock(); + } + })(); + work.catch(() => {}); + merging.push(work); + }, + }; + return { + writer, + async flush() { + try { + let cursor = 0; + while (cursor < merging.length) { + const batch = merging.slice(cursor); + cursor = merging.length; + await Promise.all(batch); + } + } finally { + await raw.close(); + } + }, + }; +} diff --git a/packages/trigger-sdk/src/v3/shared.ts b/packages/trigger-sdk/src/v3/shared.ts index abccc324dbe..2e6c635396e 100644 --- a/packages/trigger-sdk/src/v3/shared.ts +++ b/packages/trigger-sdk/src/v3/shared.ts @@ -89,9 +89,13 @@ import { type TriggerAndWaitOptions, type TriggerApiRequestOptions, type TriggerOptions, + type TaskConcurrency, + type ConcurrencyLimitOptions, + type ConcurrencyLimit, } from "@trigger.dev/core/v3"; import { resolveTriggerExternalDeploymentId, scopedEnvVar } from "./externalDeploymentId.js"; import { tracer } from "./tracer.js"; +import { triggerConcurrencyBody, validateConcurrencyLimitName } from "./concurrency-shared.js"; export type { AnyRunHandle, @@ -120,6 +124,82 @@ export { SubtaskUnwrapError, TaskRunPromise }; export type Context = TaskRunContext; +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 }; +} + export function queue(options: QueueOptions): Queue { resourceCatalog.registerQueueMetadata(options); @@ -129,6 +209,33 @@ 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) => {}, + * }); + * ``` + */ +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, @@ -158,6 +265,8 @@ export function createTask< | TaskOptions | TaskOptionsWithSchema ): Task | Task { + const normalizedQueue = normalizeTaskQueue(params.queue); + const task: Task = { id: params.id, description: params.description, @@ -169,7 +278,7 @@ export function createTask< payload, undefined, { - queue: params.queue?.name, + queue: normalizedQueue.queue?.name, ...options, } ); @@ -182,7 +291,7 @@ export function createTask< options, undefined, undefined, - params.queue?.name + normalizedQueue.queue?.name ); }, triggerAndWait: (payload, options, requestOptions) => { @@ -193,7 +302,7 @@ export function createTask< payload, undefined, { - queue: params.queue?.name, + queue: normalizedQueue.queue?.name, ...options, }, requestOptions @@ -214,7 +323,7 @@ export function createTask< payload, undefined, { - queue: params.queue?.name, + queue: normalizedQueue.queue?.name, ...options, } ) @@ -234,7 +343,7 @@ export function createTask< undefined, options, undefined, - params.queue?.name + normalizedQueue.queue?.name ); }, }; @@ -244,7 +353,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, @@ -257,9 +367,16 @@ export function createTask< }, }); - 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, @@ -312,6 +429,8 @@ export function createSchemaTask< ? getSchemaParseFn>(params.schema) : undefined; + const normalizedQueue = normalizeTaskQueue(params.queue); + const task: TaskWithSchema = { id: params.id, description: params.description, @@ -323,7 +442,7 @@ export function createSchemaTask< payload, parsePayload, { - queue: params.queue?.name, + queue: normalizedQueue.queue?.name, ...options, }, requestOptions @@ -337,7 +456,7 @@ export function createSchemaTask< options, parsePayload, requestOptions, - params.queue?.name + normalizedQueue.queue?.name ); }, triggerAndWait: (payload, options) => { @@ -348,7 +467,7 @@ export function createSchemaTask< payload, parsePayload, { - queue: params.queue?.name, + queue: normalizedQueue.queue?.name, ...options, } ) @@ -368,7 +487,7 @@ export function createSchemaTask< payload, parsePayload, { - queue: params.queue?.name, + queue: normalizedQueue.queue?.name, ...options, } ) @@ -388,7 +507,7 @@ export function createSchemaTask< parsePayload, options, undefined, - params.queue?.name + normalizedQueue.queue?.name ); }, }; @@ -398,7 +517,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, @@ -412,9 +532,16 @@ 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, @@ -715,7 +842,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, @@ -975,7 +1103,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, @@ -1240,7 +1369,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, @@ -1505,7 +1635,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, @@ -2047,7 +2179,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, @@ -2097,7 +2230,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, @@ -2264,11 +2396,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, @@ -2317,7 +2446,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, @@ -2402,11 +2532,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, @@ -2585,7 +2712,8 @@ async function triggerAndWait_internal; +}; + +/** + * AI SDK 5/6 rebuild input from the original prompt on every step. Retain an + * override and append only the newly completed response messages. The same + * logic works when the installed AI SDK already retains overrides: its input + * is not concatenated again. State belongs to one generation, not one turn. + */ +export function retainStepMessages( + prepare: (input: T) => Promise +): (input: T) => Promise { + let retained: ModelMessage[] | undefined; + let responseCount = 0; + let previousSteps = -1; + return async (input) => { + if (input.steps.length <= previousSteps) { + retained = undefined; + responseCount = 0; + } + const responses = input.steps.at(-1)?.response.messages ?? []; + const messages = retained ? [...retained, ...responses.slice(responseCount)] : input.messages; + const result = await prepare({ ...input, messages }); + retained = result?.messages ?? (retained ? messages : undefined); + responseCount = responses.length; + previousSteps = input.steps.length; + return retained ? ({ ...result, messages: retained } as R) : result; + }; +} + +export function steeringMarkers(messages: readonly UIMessage[]) { + return messages.flatMap((message) => + message.parts.flatMap((part) => { + if (part.type !== PENDING_MESSAGE_INJECTED_TYPE) return []; + const p = part as { + id?: string; + data?: { messageIds?: string[]; messages?: Array<{ id: string; text: string }> }; + }; + return typeof p.id === "string" && Array.isArray(p.data?.messageIds) + ? [{ id: p.id, messageIds: p.data.messageIds, messages: p.data.messages ?? [] }] + : []; + }) + ); +} + +/** + * Project an inline UI event into a real model-message boundary. The public + * assistant keeps its identity; the model sees complete preceding tool steps, + * then the exact prepared input, then the subsequent assistant steps. + */ +export async function convertSteeredMessages( + messages: UIMessage[], + convert: (messages: UIMessage[]) => Promise, + injections: ReadonlyMap, + context: readonly UIMessage[] = messages, + requireInputPresence = false +): Promise { + const markers = steeringMarkers([...context, ...messages]); + if (markers.length === 0 && injections.size === 0) return convert(messages); + const inlineIds = new Set(markers.flatMap((marker) => marker.messageIds)); + const uiById = new Map([...context, ...messages].map((message) => [message.id, message])); + const out: ModelMessage[] = []; + const standalone = new Map( + [...injections.values()].flatMap((entry) => entry.messageIds.map((id) => [id, entry] as const)) + ); + const emitted = new Set(); + let batch: UIMessage[] = []; + const flush = async () => { + if (batch.length) out.push(...(await convert(batch))); + batch = []; + }; + for (const message of messages) { + if (inlineIds.has(message.id)) continue; + const preparedInput = standalone.get(message.id); + if (preparedInput) { + await flush(); + if (!emitted.has(preparedInput.id)) out.push(...preparedInput.messages); + emitted.add(preparedInput.id); + continue; + } + if (message.role !== "assistant") { + batch.push(message); + continue; + } + let parts: UIMessage["parts"] = []; + for (const part of message.parts) { + const marker = steeringMarkers([{ ...message, parts: [part] }])[0]; + if (!marker) { + parts.push(part); + continue; + } + if (parts.length) batch.push({ ...message, parts }); + parts = []; + await flush(); + // An explicit history edit removing input is authoritative. The UI + // marker can still describe what happened, without resurrecting it in + // future model context from its fallback text or private prepared form. + const prepared = injections.get(marker.id); + if ( + requireInputPresence && + !marker.messageIds.every( + (id) => uiById.has(id) || prepared?.trimmedMessageIds?.includes(id) + ) + ) + continue; + if (prepared) out.push(...prepared.messages); + else { + // Old/custom stores may have only the UI transcript. Recover ordinary + // user input, never expose the private model forms in the marker. + const users = marker.messageIds.flatMap((id) => { + const ui = uiById.get(id); + if (ui) return [ui]; + const text = marker.messages.find((m) => m.id === id)?.text; + return text === undefined + ? [] + : [{ id, role: "user" as const, parts: [{ type: "text" as const, text }] }]; + }); + if (users.length) out.push(...(await convert(users))); + } + } + if (parts.length) batch.push({ ...message, parts }); + } + await flush(); + return out; +} diff --git a/packages/trigger-sdk/src/v3/transcriptStorage.ts b/packages/trigger-sdk/src/v3/transcriptStorage.ts index 2f102a1fe1a..1b25f19334f 100644 --- a/packages/trigger-sdk/src/v3/transcriptStorage.ts +++ b/packages/trigger-sdk/src/v3/transcriptStorage.ts @@ -6,6 +6,7 @@ import { } from "@trigger.dev/core/v3"; import type { ModelMessage, UIMessage } from "ai"; import { readChatSnapshot, writeChatSnapshot } from "./chatSnapshotIo.js"; +import { steeringMarkers, type SteeringInjection } from "./steeringContext.js"; /** * One change to a transcript. Ops address messages by id; position is the @@ -328,7 +329,8 @@ export function snapshotTranscriptStorage(): TranscriptStorage { messages: entries, state: normalizeRuntimeStateForWindow( changeset.transcript.state, - entries.map((entry) => entry.id) + entries.map((entry) => entry.message), + changeset.transcript.entries.map((entry) => entry.id) ), lastOutEventId: changeset.cursors?.lastOutEventId, lastInEventId: changeset.cursors?.lastInEventId, @@ -443,6 +445,8 @@ export type TranscriptRuntimeState = { v: 1; compaction?: { modelMessages: ModelMessage[]; throughId: string }; injections?: ModelLaneInjection[]; + /** Prepared steering model forms; boundaries live in the corresponding UI markers. */ + steering?: SteeringInjection[]; /** `chat.inject` messages queued but not yet drained into a turn when the save happened. */ queued?: ModelMessage[]; }; @@ -494,12 +498,20 @@ export function trimTranscriptForSnapshot( */ export function normalizeRuntimeStateForWindow( state: unknown, - windowIds: Iterable + window: Iterable, + originalIds: Iterable = [] ): unknown { const runtimeState = parseTranscriptRuntimeState(state); if (!runtimeState) return state; - const ids = new Set(windowIds); + const retained = [...window]; + const ids = new Set(retained.map((entry) => (typeof entry === "string" ? entry : entry.id))); + const markerIds = new Set( + steeringMarkers(retained.filter((entry): entry is UIMessage => typeof entry !== "string")).map( + (marker) => marker.id + ) + ); + const previousIds = new Set(originalIds); const compaction = runtimeState.compaction; const watermarkOutside = compaction !== undefined && compaction.throughId !== "" && !ids.has(compaction.throughId); @@ -508,7 +520,31 @@ export function normalizeRuntimeStateForWindow( (injection) => injection.afterId === "" || ids.has(injection.afterId) ); - if (!watermarkOutside && injections?.length === runtimeState.injections?.length) { + let steeringChanged = false; + const steering = runtimeState.steering?.flatMap((entry) => { + if (!markerIds.has(entry.id) && !entry.messageIds.some((id) => ids.has(id))) { + steeringChanged = true; + return []; + } + // Preserve prepared input when its boundary survives the window cutoff. + // Record only IDs actually trimmed from this save, so an earlier explicit + // history deletion cannot become a retention exception on the next boot. + const trimmedMessageIds = entry.messageIds.filter( + (id) => !ids.has(id) && (previousIds.has(id) || entry.trimmedMessageIds?.includes(id)) + ); + if ( + trimmedMessageIds.length === (entry.trimmedMessageIds?.length ?? 0) && + trimmedMessageIds.every((id, index) => id === entry.trimmedMessageIds?.[index]) + ) + return [entry]; + steeringChanged = true; + return [{ ...entry, trimmedMessageIds }]; + }); + if ( + !watermarkOutside && + injections?.length === runtimeState.injections?.length && + !steeringChanged + ) { return state; } @@ -523,6 +559,7 @@ export function normalizeRuntimeStateForWindow( } : {}), ...(injections ? { injections } : {}), + ...(steering ? { steering } : {}), }; } @@ -551,6 +588,32 @@ export function parseTranscriptRuntimeState(value: unknown): TranscriptRuntimeSt : []; }); } + if (Array.isArray(record.steering)) { + out.steering = record.steering.flatMap((entry: unknown) => { + if (!entry || typeof entry !== "object") return []; + const value = entry as Record; + return typeof value.id === "string" && + Array.isArray(value.messageIds) && + value.messageIds.every((id) => typeof id === "string") && + Array.isArray(value.messages) + ? [ + { + id: value.id, + messageIds: value.messageIds as string[], + messages: value.messages as ModelMessage[], + ...(Array.isArray(value.trimmedMessageIds) + ? { + trimmedMessageIds: value.trimmedMessageIds.filter( + (id): id is string => + typeof id === "string" && (value.messageIds as string[]).includes(id) + ), + } + : {}), + }, + ] + : []; + }); + } if (Array.isArray(record.queued) && record.queued.length > 0) { out.queued = record.queued as ModelMessage[]; } diff --git a/packages/trigger-sdk/test/steering-compaction-lanes.test.ts b/packages/trigger-sdk/test/steering-compaction-lanes.test.ts index a075f80a2f4..e107ee06be4 100644 --- a/packages/trigger-sdk/test/steering-compaction-lanes.test.ts +++ b/packages/trigger-sdk/test/steering-compaction-lanes.test.ts @@ -189,7 +189,18 @@ describe("chat.agent steering with compaction in the same turn", () => { expect(completedTurns[0]?.newMessages).toContain("COMPACTED_TOOL_SENTINEL"); expect(completedTurns[0]?.newMessages).toContain('"toolCallId":"tc-1"'); expect(completedTurns[0]?.newMessages).toContain("done"); - expect(completedTurns[0]?.newMessages).toEqual(completedTurns[0]?.fullDelta); + // UI storage keeps one assistant; model deltas split it at the durable + // injection marker. The same content now has its actual causal order. + const delta = JSON.parse(completedTurns[0]!.newMessages); + expect(delta.map((m: { role: string }) => m.role)).toEqual([ + "user", + "assistant", + "tool", + "user", + "assistant", + ]); + expect(delta).toHaveLength(JSON.parse(completedTurns[0]!.fullDelta).length); + expect(JSON.stringify(delta[3])).toContain("steer-me"); } finally { toolGate.resolve(); await harness.close();