use new plan type on runs as fallback during dequeue

nicktrn · nicktrn · commit 5d75706dcc07 · 2025-08-14T10:32:50.000+01:00
diff --git a/apps/webapp/app/v3/runEngine.server.ts b/apps/webapp/app/v3/runEngine.server.ts
@@ -110,14 +110,23 @@ function createRunEngine() {
         const plan = await getCurrentPlan(orgId);
 
         if (!plan) {
-          return { isPaying: false };
+          return {
+            isPaying: false,
+            type: "free",
+          };
         }
 
         if (!plan.v3Subscription) {
-          return { isPaying: false };
+          return {
+            isPaying: false,
+            type: "free",
+          };
         }
 
-        return { isPaying: plan.v3Subscription.isPaying };
+        return {
+          isPaying: plan.v3Subscription.isPaying,
+          type: plan.v3Subscription.plan?.type ?? "free",
+        };
       },
     },
   });
diff --git a/internal-packages/cache/src/index.ts b/internal-packages/cache/src/index.ts
@@ -3,6 +3,8 @@ export {
   DefaultStatefulContext,
   Namespace,
   type Cache as UnkeyCache,
+  type CacheError,
 } from "@unkey/cache";
+export { type Result, Ok, Err } from "@unkey/error";
 export { MemoryStore } from "@unkey/cache/stores";
 export { RedisCacheStore } from "./stores/redis.js";
diff --git a/internal-packages/database/prisma/migrations/20250814092224_add_task_run_plan_type/migration.sql b/internal-packages/database/prisma/migrations/20250814092224_add_task_run_plan_type/migration.sql
@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "TaskRun" ADD COLUMN     "planType" TEXT;
diff --git a/internal-packages/database/prisma/schema.prisma b/internal-packages/database/prisma/schema.prisma
@@ -707,6 +707,9 @@ model TaskRun {
   /// Run error
   error Json?
 
+  /// Organization's billing plan type (cached for fallback when billing API fails)
+  planType String?
+
   maxDurationInSeconds Int?
 
   @@unique([oneTimeUseToken])
diff --git a/internal-packages/run-engine/src/engine/billingCache.ts b/internal-packages/run-engine/src/engine/billingCache.ts
@@ -0,0 +1,92 @@
+import {
+  createCache,
+  DefaultStatefulContext,
+  MemoryStore,
+  Namespace,
+  Ok,
+  RedisCacheStore,
+  type UnkeyCache,
+  type CacheError,
+  type Result,
+} from "@internal/cache";
+import type { RedisOptions } from "@internal/redis";
+import type { Logger } from "@trigger.dev/core/logger";
+import type { RunEngineOptions } from "./types.js";
+
+// Cache TTLs for billing information - shorter than other caches since billing can change
+const BILLING_FRESH_TTL = 60000 * 5; // 5 minutes
+const BILLING_STALE_TTL = 60000 * 10; // 10 minutes
+
+export type BillingPlan = {
+  isPaying: boolean;
+  type: "free" | "paid" | "enterprise";
+};
+
+export type BillingCacheOptions = {
+  billingOptions?: RunEngineOptions["billing"];
+  redisOptions: RedisOptions;
+  logger: Logger;
+};
+
+export class BillingCache {
+  private readonly cache: UnkeyCache<{
+    currentPlan: BillingPlan;
+  }>;
+  private readonly logger: Logger;
+  private readonly billingOptions?: RunEngineOptions["billing"];
+
+  constructor(options: BillingCacheOptions) {
+    this.logger = options.logger;
+    this.billingOptions = options.billingOptions;
+
+    // Initialize cache
+    const ctx = new DefaultStatefulContext();
+    const memory = new MemoryStore({ persistentMap: new Map() });
+    const redisCacheStore = new RedisCacheStore({
+      name: "billing-cache",
+      connection: {
+        ...options.redisOptions,
+        keyPrefix: "engine:billing:cache:",
+      },
+      useModernCacheKeyBuilder: true,
+    });
+
+    this.cache = createCache({
+      currentPlan: new Namespace<BillingPlan>(ctx, {
+        stores: [memory, redisCacheStore],
+        fresh: BILLING_FRESH_TTL,
+        stale: BILLING_STALE_TTL,
+      }),
+    });
+  }
+
+  /**
+   * Gets the current billing plan for an organization
+   * Returns a Result that allows the caller to handle errors and missing values
+   */
+  async getCurrentPlan(orgId: string): Promise<Result<BillingPlan | undefined, CacheError>> {
+    if (!this.billingOptions?.getCurrentPlan) {
+      // Return a successful result with default free plan
+      return Ok({ isPaying: false, type: "free" });
+    }
+
+    return await this.cache.currentPlan.swr(orgId, async () => {
+      // This is safe because options can't change at runtime
+      const planResult = await this.billingOptions!.getCurrentPlan(orgId);
+      return { isPaying: planResult.isPaying, type: planResult.type };
+    });
+  }
+
+  /**
+   * Invalidates the billing cache for an organization when their plan changes
+   * Runs in background and handles all errors internally
+   */
+  invalidate(orgId: string): void {
+    this.cache.currentPlan.remove(orgId).catch((error) => {
+      this.logger.warn("Failed to invalidate billing cache", {
+        orgId,
+        error: error instanceof Error ? error.message : String(error),
+      });
+    });
+  }
+}
diff --git a/internal-packages/run-engine/src/engine/index.ts b/internal-packages/run-engine/src/engine/index.ts
@@ -1,3 +1,4 @@
+import { BillingCache } from "./billingCache.js";
 import { createRedisClient, Redis } from "@internal/redis";
 import { getMeter, Meter, startSpan, trace, Tracer } from "@internal/tracing";
 import { Logger } from "@trigger.dev/core/logger";
@@ -78,6 +79,8 @@ export class RunEngine {
   pendingVersionSystem: PendingVersionSystem;
   raceSimulationSystem: RaceSimulationSystem = new RaceSimulationSystem();
 
+  private readonly billingCache: BillingCache;
+
   constructor(private readonly options: RunEngineOptions) {
     this.logger = options.logger ?? new Logger("RunEngine", this.options.logLevel ?? "info");
     this.prisma = options.prisma;
@@ -292,12 +295,18 @@ export class RunEngine {
       redisOptions: this.options.cache?.redis ?? this.options.runLock.redis,
     });
 
+    this.billingCache = new BillingCache({
+      billingOptions: this.options.billing,
+      redisOptions: this.options.cache?.redis ?? this.options.runLock.redis,
+      logger: this.logger,
+    });
+
     this.dequeueSystem = new DequeueSystem({
       resources,
       executionSnapshotSystem: this.executionSnapshotSystem,
       runAttemptSystem: this.runAttemptSystem,
       machines: this.options.machines,
-      billing: this.options.billing,
+      billingCache: this.billingCache,
       redisOptions: this.options.cache?.redis ?? this.options.runLock.redis,
     });
   }
@@ -366,13 +375,37 @@ export class RunEngine {
       "trigger",
       async (span) => {
         const status = delayUntil ? "DELAYED" : "PENDING";
+        const runId = RunId.fromFriendlyId(friendlyId);
+
+        // Get billing information to store with the run
+        let planType: string | undefined;
+        const currentPlan = await this.billingCache.getCurrentPlan(environment.organization.id);
+
+        if (currentPlan.err || !currentPlan.val) {
+          // If billing lookup fails, don't block the trigger - planType will be null
+          this.logger.warn(
+            "Failed to get billing info during trigger, proceeding without planType",
+            {
+              orgId: environment.organization.id,
+              runId,
+              error:
+                currentPlan.err instanceof Error
+                  ? currentPlan.err.message
+                  : String(currentPlan.err),
+              hasValue: !!currentPlan.val,
+            }
+          );
+          planType = undefined;
+        } else {
+          planType = currentPlan.val.type;
+        }
 
         //create run
         let taskRun: TaskRun;
         try {
           taskRun = await prisma.taskRun.create({
             data: {
-              id: RunId.fromFriendlyId(friendlyId),
+              id: runId,
               engine: "V2",
               status,
               number,
@@ -431,6 +464,7 @@ export class RunEngine {
               scheduleInstanceId,
               createdAt,
               bulkActionGroupIds: bulkActionId ? [bulkActionId] : undefined,
+              planType,
               executionSnapshots: {
                 create: {
                   engine: "V2",
@@ -1355,6 +1389,6 @@ export class RunEngine {
    * Runs in background and handles all errors internally
    */
   invalidateBillingCache(orgId: string): void {
-    this.dequeueSystem.invalidateBillingCache(orgId);
+    this.billingCache.invalidate(orgId);
   }
 }
diff --git a/internal-packages/run-engine/src/engine/systems/dequeueSystem.ts b/internal-packages/run-engine/src/engine/systems/dequeueSystem.ts
@@ -1,12 +1,5 @@
-import {
-  createCache,
-  DefaultStatefulContext,
-  MemoryStore,
-  Namespace,
-  RedisCacheStore,
-  UnkeyCache,
-} from "@internal/cache";
 import type { RedisOptions } from "@internal/redis";
+import type { BillingCache } from "../billingCache.js";
 import { startSpan } from "@internal/tracing";
 import { assertExhaustive } from "@trigger.dev/core";
 import { DequeuedMessage, RetryOptions } from "@trigger.dev/core/v3";
@@ -20,66 +13,25 @@ import { RunEngineOptions } from "../types.js";
 import { ExecutionSnapshotSystem, getLatestExecutionSnapshot } from "./executionSnapshotSystem.js";
 import { RunAttemptSystem } from "./runAttemptSystem.js";
 import { SystemResources } from "./systems.js";
-import { ServiceValidationError } from "../errors.js";
 
 export type DequeueSystemOptions = {
   resources: SystemResources;
   machines: RunEngineOptions["machines"];
   executionSnapshotSystem: ExecutionSnapshotSystem;
   runAttemptSystem: RunAttemptSystem;
-  billing?: RunEngineOptions["billing"];
+  billingCache: BillingCache;
   redisOptions: RedisOptions;
 };
 
-// Cache TTLs for billing information - shorter than other caches since billing can change
-const BILLING_FRESH_TTL = 60000 * 5; // 5 minutes
-const BILLING_STALE_TTL = 60000 * 10; // 10 minutes
-
 export class DequeueSystem {
   private readonly $: SystemResources;
   private readonly executionSnapshotSystem: ExecutionSnapshotSystem;
   private readonly runAttemptSystem: RunAttemptSystem;
-  private readonly billingCache: UnkeyCache<{
-    billing: { isPaying: boolean };
-  }>;
 
   constructor(private readonly options: DequeueSystemOptions) {
     this.$ = options.resources;
     this.executionSnapshotSystem = options.executionSnapshotSystem;
     this.runAttemptSystem = options.runAttemptSystem;
-
-    // Initialize billing cache
-    const ctx = new DefaultStatefulContext();
-    const memory = new MemoryStore({ persistentMap: new Map() });
-    const redisCacheStore = new RedisCacheStore({
-      name: "dequeue-system",
-      connection: {
-        ...options.redisOptions,
-        keyPrefix: "engine:dequeue-system:cache:",
-      },
-      useModernCacheKeyBuilder: true,
-    });
-
-    this.billingCache = createCache({
-      billing: new Namespace<{ isPaying: boolean }>(ctx, {
-        stores: [memory, redisCacheStore],
-        fresh: BILLING_FRESH_TTL,
-        stale: BILLING_STALE_TTL,
-      }),
-    });
-  }
-
-  /**
-   * Invalidates the billing cache for an organization when their plan changes
-   * Runs in background and handles all errors internally
-   */
-  invalidateBillingCache(orgId: string): void {
-    this.billingCache.billing.remove(orgId).catch((error) => {
-      this.$.logger.warn("Failed to invalidate billing cache", {
-        orgId,
-        error: error.message,
-      });
-    });
   }
 
   /**
@@ -432,8 +384,29 @@ export class DequeueSystem {
               const currentAttemptNumber = lockedTaskRun.attemptNumber ?? 0;
               const nextAttemptNumber = currentAttemptNumber + 1;
 
-              // Get billing information if available
-              const billing = await this.#getBillingInfo({ orgId, runId });
+              // Get billing information if available, with fallback to TaskRun.planType
+              const billingResult = await this.options.billingCache.getCurrentPlan(orgId);
+
+              let isPaying: boolean;
+              if (billingResult.err || !billingResult.val) {
+                // Fallback to stored planType on TaskRun if billing cache fails or returns no value
+                this.$.logger.warn(
+                  "Billing cache failed or returned no value, falling back to TaskRun.planType",
+                  {
+                    orgId,
+                    runId,
+                    error:
+                      billingResult.err instanceof Error
+                        ? billingResult.err.message
+                        : String(billingResult.err),
+                    currentPlan: billingResult.val,
+                  }
+                );
+
+                isPaying = lockedTaskRun.planType !== null && lockedTaskRun.planType !== "free";
+              } else {
+                isPaying = billingResult.val.isPaying;
+              }
 
               const newSnapshot = await this.executionSnapshotSystem.createExecutionSnapshot(
                 prisma,
@@ -503,7 +476,11 @@ export class DequeueSystem {
                 project: {
                   id: lockedTaskRun.projectId,
                 },
-                billing,
+                billing: {
+                  currentPlan: {
+                    isPaying,
+                  },
+                },
               } satisfies DequeuedMessage;
             }
           );
@@ -668,37 +645,4 @@ export class DequeueSystem {
       });
     });
   }
-
-  async #getBillingInfo({
-    orgId,
-    runId,
-  }: {
-    orgId: string;
-    runId: string;
-  }): Promise<{ currentPlan: { isPaying: boolean } }> {
-    if (!this.options.billing?.getCurrentPlan) {
-      return { currentPlan: { isPaying: false } };
-    }
-
-    const result = await this.billingCache.billing.swr(orgId, async () => {
-      // This is safe because options can't change at runtime
-      const planResult = await this.options.billing!.getCurrentPlan(orgId);
-
-      return { isPaying: planResult.isPaying };
-    });
-
-    if (result.err) {
-      throw result.err;
-    }
-
-    if (!result.val) {
-      throw new ServiceValidationError(
-        `Could not resolve billing information for organization ${orgId}`,
-        undefined,
-        { orgId, runId }
-      );
-    }
-
-    return { currentPlan: { isPaying: result.val.isPaying } };
-  }
 }
diff --git a/internal-packages/run-engine/src/engine/types.ts b/internal-packages/run-engine/src/engine/types.ts

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+-- AlterTable`
	`2`	`+ALTER TABLE "TaskRun" ADD COLUMN "planType" TEXT;`