start using new entitlement response

Author: nicktrn

apps/webapp/app/runEngine/services/triggerTask.server.ts

@@ -123,13 +123,35 @@ export class RunEngineTriggerTaskService {
         throw tagValidation.error;
       }
 
-      // Validate entitlement
-      const entitlementValidation = await this.validator.validateEntitlement({
-        environment,
-      });
+      // Validate entitlement (unless skipChecks is enabled)
+      let planType: string | undefined;
+
+      if (!options.skipChecks) {
+        const entitlementValidation = await this.validator.validateEntitlement({
+          environment,
+        });
 
-      if (!entitlementValidation.ok) {
-        throw entitlementValidation.error;
+        if (!entitlementValidation.ok) {
+          throw entitlementValidation.error;
+        }
+
+        // Extract plan type from entitlement response
+        planType = entitlementValidation.plan?.type;
+      } else {
+        // When skipChecks is enabled, planType should be passed via options
+        planType = options.planType;
+
+        if (!planType) {
+          logger.warn("Plan type not set but skipChecks is enabled", {
+            taskId,
+            environment: {
+              id: environment.id,
+              type: environment.type,
+              projectId: environment.projectId,
+              organizationId: environment.organizationId,
+            },
+          });
+        }
       }
 
       const [parseDelayError, delayUntil] = await tryCatch(parseDelay(body.options?.delay));
@@ -313,6 +335,7 @@ export class RunEngineTriggerTaskService {
                   scheduleInstanceId: options.scheduleInstanceId,
                   createdAt: options.overrideCreatedAt,
                   bulkActionId: body.options?.bulkActionId,
+                  planType,
                 },
                 this.prisma
               );

apps/webapp/app/runEngine/types.ts

@@ -1,12 +1,7 @@
-import { BackgroundWorker, TaskRun } from "@trigger.dev/database";
-
-import {
-  IOPacket,
-  RunChainState,
-  TaskRunError,
-  TriggerTaskRequestBody,
-} from "@trigger.dev/core/v3";
-import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
+import type { BackgroundWorker, TaskRun } from "@trigger.dev/database";
+import type { IOPacket, TaskRunError, TriggerTaskRequestBody } from "@trigger.dev/core/v3";
+import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
+import type { ReportUsagePlan } from "@trigger.dev/platform";
 
 export type TriggerTaskServiceOptions = {
   idempotencyKey?: string;
@@ -22,6 +17,7 @@ export type TriggerTaskServiceOptions = {
   skipChecks?: boolean;
   oneTimeUseToken?: string;
   overrideCreatedAt?: Date;
+  planType?: string;
 };
 
 // domain/triggerTask.ts
@@ -112,9 +108,19 @@ export type ValidationResult =
       error: Error;
     };
 
+export type EntitlementValidationResult =
+  | {
+      ok: true;
+      plan?: ReportUsagePlan;
+    }
+  | {
+      ok: false;
+      error: Error;
+    };
+
 export interface TriggerTaskValidator {
   validateTags(params: TagValidationParams): ValidationResult;
-  validateEntitlement(params: EntitlementValidationParams): Promise<ValidationResult>;
+  validateEntitlement(params: EntitlementValidationParams): Promise<EntitlementValidationResult>;
   validateMaxAttempts(params: MaxAttemptsValidationParams): ValidationResult;
   validateParentRun(params: ParentRunValidationParams): ValidationResult;
 }

apps/webapp/app/runEngine/validators/triggerTaskValidator.ts

@@ -4,8 +4,9 @@ import { getEntitlement } from "~/services/platform.v3.server";
 import { MAX_ATTEMPTS, OutOfEntitlementError } from "~/v3/services/triggerTask.server";
 import { isFinalRunStatus } from "~/v3/taskStatus";
 import { EngineServiceValidationError } from "../concerns/errors";
-import {
+import type {
   EntitlementValidationParams,
+  EntitlementValidationResult,
   MaxAttemptsValidationParams,
   ParentRunValidationParams,
   TagValidationParams,
@@ -37,7 +38,9 @@ export class DefaultTriggerTaskValidator implements TriggerTaskValidator {
     return { ok: true };
   }
 
-  async validateEntitlement(params: EntitlementValidationParams): Promise<ValidationResult> {
+  async validateEntitlement(
+    params: EntitlementValidationParams
+  ): Promise<EntitlementValidationResult> {
     const { environment } = params;
 
     if (environment.type === "DEVELOPMENT") {
@@ -53,7 +56,7 @@ export class DefaultTriggerTaskValidator implements TriggerTaskValidator {
       };
     }
 
-    return { ok: true };
+    return { ok: true, plan: result?.plan };
   }
 
   validateMaxAttempts(params: MaxAttemptsValidationParams): ValidationResult {

apps/webapp/app/services/platform.v3.server.ts

@@ -10,6 +10,8 @@ import {
   type MachineCode,
   type UpdateBillingAlertsRequest,
   type BillingAlertsResult,
+  type ReportUsageResult,
+  type ReportUsagePlan,
 } from "@trigger.dev/platform";
 import { createCache, DefaultStatefulContext, Namespace } from "@unkey/cache";
 import { MemoryStore } from "@unkey/cache/stores";
@@ -432,7 +434,9 @@ export async function reportComputeUsage(request: Request) {
   });
 }
 
-export async function getEntitlement(organizationId: string) {
+export async function getEntitlement(
+  organizationId: string
+): Promise<ReportUsageResult | undefined> {
   if (!client) return undefined;
 
   try {

apps/webapp/app/v3/services/triggerTask.server.ts

@@ -33,6 +33,7 @@ export type TriggerTaskServiceOptions = {
   queueTimestamp?: Date;
   overrideCreatedAt?: Date;
   replayedFromTaskRunFriendlyId?: string;
+  planType?: string;
 };
 
 export class OutOfEntitlementError extends Error {