add billing info to dequeued message w/o cache

Author: nicktrn

apps/supervisor/src/index.ts

@@ -247,6 +247,7 @@ class ManagedSupervisor {
           nextAttemptNumber: message.run.attemptNumber,
           snapshotId: message.snapshot.id,
           snapshotFriendlyId: message.snapshot.friendlyId,
+          isPaidTier: message.billing?.currentPlan.isPaying ?? false,
         });
 
         // Disabled for now

apps/webapp/app/v3/runEngine.server.ts

@@ -1,7 +1,7 @@
 import { RunEngine } from "@internal/run-engine";
 import { $replica, prisma } from "~/db.server";
 import { env } from "~/env.server";
-import { defaultMachine } from "~/services/platform.v3.server";
+import { defaultMachine, getCurrentPlan } from "~/services/platform.v3.server";
 import { singleton } from "~/utils/singleton";
 import { allMachines } from "./machinePresets.server";
 import { meter, tracer } from "./tracer.server";
@@ -105,6 +105,21 @@ function createRunEngine() {
       SUSPENDED: env.RUN_ENGINE_TIMEOUT_SUSPENDED,
     },
     retryWarmStartThresholdMs: env.RUN_ENGINE_RETRY_WARM_START_THRESHOLD_MS,
+    billing: {
+      getCurrentPlan: async (orgId: string) => {
+        const plan = await getCurrentPlan(orgId);
+
+        if (!plan) {
+          return { isPaying: false };
+        }
+
+        if (!plan.v3Subscription) {
+          return { isPaying: false };
+        }
+
+        return { isPaying: plan.v3Subscription.isPaying };
+      },
+    },
   });
 
   return engine;

internal-packages/run-engine/src/engine/index.ts

@@ -297,6 +297,7 @@ export class RunEngine {
       executionSnapshotSystem: this.executionSnapshotSystem,
       runAttemptSystem: this.runAttemptSystem,
       machines: this.options.machines,
+      billing: this.options.billing,
     });
   }
 

internal-packages/run-engine/src/engine/systems/dequeueSystem.ts

@@ -1,5 +1,5 @@
 import { startSpan } from "@internal/tracing";
-import { assertExhaustive } from "@trigger.dev/core";
+import { assertExhaustive, tryCatch } from "@trigger.dev/core";
 import { DequeuedMessage, RetryOptions } from "@trigger.dev/core/v3";
 import { getMaxDuration } from "@trigger.dev/core/v3/isomorphic";
 import { PrismaClientOrTransaction } from "@trigger.dev/database";
@@ -17,6 +17,7 @@ export type DequeueSystemOptions = {
   machines: RunEngineOptions["machines"];
   executionSnapshotSystem: ExecutionSnapshotSystem;
   runAttemptSystem: RunAttemptSystem;
+  billing?: RunEngineOptions["billing"];
 };
 
 export class DequeueSystem {
@@ -380,6 +381,24 @@ export class DequeueSystem {
               const currentAttemptNumber = lockedTaskRun.attemptNumber ?? 0;
               const nextAttemptNumber = currentAttemptNumber + 1;
 
+              // Get billing information if available
+              let isPaying = false;
+              if (this.options.billing?.getCurrentPlan) {
+                const [error, planResult] = await tryCatch(
+                  this.options.billing.getCurrentPlan(orgId)
+                );
+
+                if (error) {
+                  this.$.logger.error("Failed to get billing information", {
+                    orgId,
+                    runId,
+                    error: error.message,
+                  });
+                } else {
+                  isPaying = planResult.isPaying;
+                }
+              }
+
               const newSnapshot = await this.executionSnapshotSystem.createExecutionSnapshot(
                 prisma,
                 {
@@ -448,6 +467,11 @@ export class DequeueSystem {
                 project: {
                   id: lockedTaskRun.projectId,
                 },
+                billing: {
+                  currentPlan: {
+                    isPaying,
+                  },
+                },
               } satisfies DequeuedMessage;
             }
           );

internal-packages/run-engine/src/engine/types.ts

@@ -30,6 +30,9 @@ export type RunEngineOptions = {
     machines: Record<string, MachinePreset>;
     baseCostInCents: number;
   };
+  billing?: {
+    getCurrentPlan: (orgId: string) => Promise<{ isPaying: boolean }>;
+  };
   queue: {
     redis: RedisOptions;
     shardCount?: number;

packages/core/src/v3/schemas/runEngine.ts

@@ -261,5 +261,12 @@ export const DequeuedMessage = z.object({
   project: z.object({
     id: z.string(),
   }),
+  billing: z
+    .object({
+      currentPlan: z.object({
+        isPaying: z.boolean(),
+      }),
+    })
+    .optional(),
 });
 export type DequeuedMessage = z.infer<typeof DequeuedMessage>;