make max dequeue count configurable from supervisor

Author: nicktrn

apps/supervisor/src/env.ts

@@ -31,6 +31,7 @@ const Env = z.object({
   // Dequeue settings (provider mode)
   TRIGGER_DEQUEUE_ENABLED: BoolEnv.default("true"),
   TRIGGER_DEQUEUE_INTERVAL_MS: z.coerce.number().int().default(1000),
+  TRIGGER_DEQUEUE_MAX_RUN_COUNT: z.coerce.number().int().default(10),
 
   // Optional services
   TRIGGER_WARM_START_URL: z.string().optional(),

apps/supervisor/src/index.ts

@@ -113,6 +113,7 @@ class ManagedSupervisor {
       managedWorkerSecret: env.MANAGED_WORKER_SECRET,
       dequeueIntervalMs: env.TRIGGER_DEQUEUE_INTERVAL_MS,
       queueConsumerEnabled: env.TRIGGER_DEQUEUE_ENABLED,
+      maxRunCount: env.TRIGGER_DEQUEUE_MAX_RUN_COUNT,
       runNotificationsEnabled: env.TRIGGER_WORKLOAD_API_ENABLED,
       preDequeue: async () => {
         if (this.isKubernetes) {

apps/webapp/app/routes/engine.v1.worker-actions.dequeue.ts

@@ -1,5 +1,8 @@
 import { json, TypedResponse } from "@remix-run/server-runtime";
-import { WorkerApiDequeueRequestBody, WorkerApiDequeueResponseBody } from "@trigger.dev/core/v3/workers";
+import {
+  WorkerApiDequeueRequestBody,
+  WorkerApiDequeueResponseBody,
+} from "@trigger.dev/core/v3/workers";
 import { createActionWorkerApiRoute } from "~/services/routeBuilders/apiBuilder.server";
 
 export const action = createActionWorkerApiRoute(
@@ -10,6 +13,7 @@ export const action = createActionWorkerApiRoute(
     return json(
       await authenticatedWorker.dequeue({
         maxResources: body.maxResources,
+        maxRunCount: body.maxRunCount,
       })
     );
   }

packages/core/src/v3/runEngineWorker/supervisor/queueConsumer.ts

@@ -7,13 +7,15 @@ type RunQueueConsumerOptions = {
   intervalMs?: number;
   preDequeue?: PreDequeueFn;
   preSkip?: PreSkipFn;
+  maxRunCount?: number;
   onDequeue: (messages: WorkerApiDequeueResponseBody) => Promise<void>;
 };
 
 export class RunQueueConsumer {
   private readonly client: SupervisorHttpClient;
   private readonly preDequeue?: PreDequeueFn;
   private readonly preSkip?: PreSkipFn;
+  private readonly maxRunCount?: number;
   private readonly onDequeue: (messages: WorkerApiDequeueResponseBody) => Promise<void>;
 
   private intervalMs: number;
@@ -24,6 +26,7 @@ export class RunQueueConsumer {
     this.intervalMs = opts.intervalMs ?? 5_000;
     this.preDequeue = opts.preDequeue;
     this.preSkip = opts.preSkip;
+    this.maxRunCount = opts.maxRunCount;
     this.onDequeue = opts.onDequeue;
     this.client = opts.client;
   }
@@ -87,6 +90,7 @@ export class RunQueueConsumer {
     try {
       const response = await this.client.dequeue({
         maxResources: preDequeueResult?.maxResources,
+        maxRunCount: this.maxRunCount,
       });
 
       if (!response.success) {

packages/core/src/v3/runEngineWorker/supervisor/schemas.ts

@@ -66,6 +66,7 @@ export type WorkerApiConnectResponseBody = z.infer<typeof WorkerApiConnectRespon
 
 export const WorkerApiDequeueRequestBody = z.object({
   maxResources: MachineResources.optional(),
+  maxRunCount: z.number().optional(),
 });
 export type WorkerApiDequeueRequestBody = z.infer<typeof WorkerApiDequeueRequestBody>;
 

packages/core/src/v3/runEngineWorker/supervisor/session.ts

@@ -17,6 +17,7 @@ type SupervisorSessionOptions = SupervisorClientCommonOptions & {
   dequeueIntervalMs?: number;
   preDequeue?: PreDequeueFn;
   preSkip?: PreSkipFn;
+  maxRunCount?: number;
 };
 
 export class SupervisorSession extends EventEmitter<WorkerEvents> {
@@ -44,6 +45,7 @@ export class SupervisorSession extends EventEmitter<WorkerEvents> {
       preSkip: opts.preSkip,
       onDequeue: this.onDequeue.bind(this),
       intervalMs: opts.dequeueIntervalMs,
+      maxRunCount: opts.maxRunCount,
     });
 
     // TODO: This should be dynamic and set by (or at least overridden by) the platform