A couple tweaks to make CodeRabbit happy and fixing the tests in CI

Author: ericallam

apps/webapp/app/runEngine/concerns/idempotencyKeys.server.ts

@@ -48,30 +48,29 @@ export class IdempotencyKeyConcern {
         });
 
         // Update the existing run to remove the idempotency key
-        await this.prisma.taskRun.update({
-          where: { id: existingRun.id },
-          data: { idempotencyKey: null },
+        await this.prisma.taskRun.updateMany({
+          where: { id: existingRun.id, idempotencyKey },
+          data: { idempotencyKey: null, idempotencyKeyExpiresAt: null },
         });
       } else {
+        const associatedWaitpoint = existingRun.associatedWaitpoint;
+        const parentRunId = request.body.options?.parentRunId;
+        const resumeParentOnCompletion = request.body.options?.resumeParentOnCompletion;
         //We're using `andWait` so we need to block the parent run with a waitpoint
-        if (
-          existingRun.associatedWaitpoint &&
-          request.body.options?.resumeParentOnCompletion &&
-          request.body.options?.parentRunId
-        ) {
+        if (associatedWaitpoint && resumeParentOnCompletion && parentRunId) {
           await this.traceEventConcern.traceIdempotentRun(
             request,
             {
               existingRun,
               idempotencyKey,
-              incomplete: existingRun.associatedWaitpoint.status === "PENDING",
-              isError: existingRun.associatedWaitpoint.outputIsError,
+              incomplete: associatedWaitpoint.status === "PENDING",
+              isError: associatedWaitpoint.outputIsError,
             },
             async (event) => {
               //block run with waitpoint
               await this.engine.blockRunWithWaitpoint({
-                runId: RunId.fromFriendlyId(request.body.options!.parentRunId!),
-                waitpoints: existingRun.associatedWaitpoint!.id,
+                runId: RunId.fromFriendlyId(parentRunId),
+                waitpoints: associatedWaitpoint.id,
                 spanIdToComplete: event.spanId,
                 batch: request.options?.batchId
                   ? {

apps/webapp/app/runEngine/concerns/runChainStates.server.ts

@@ -143,7 +143,10 @@ export class DefaultRunChainStateManager implements RunChainStateManager {
 
     const environmentConcurrency = (runChainState.concurrency?.environment ?? 0) + 1;
 
-    if (environmentConcurrency > request.environment.maximumConcurrencyLimit) {
+    if (
+      request.environment.maximumConcurrencyLimit > 0 &&
+      environmentConcurrency > request.environment.maximumConcurrencyLimit
+    ) {
       const environmentDetails = `The environment has a concurrency limit of ${request.environment.maximumConcurrencyLimit}, and the chain would require ${environmentConcurrency}`;
       throw new EngineServiceValidationError(this.createDeadlockErrorMessage(environmentDetails));
     }

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

@@ -12,14 +12,14 @@ import {
   TriggerTaskRequestBody,
 } from "@trigger.dev/core/v3";
 import { RunId, stringifyDuration } from "@trigger.dev/core/v3/isomorphic";
-import { PrismaClientOrTransaction } from "@trigger.dev/database";
+import type { PrismaClientOrTransaction } from "@trigger.dev/database";
 import { createTags } from "~/models/taskRunTag.server";
-import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
+import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
 import { parseDelay } from "~/utils/delays";
 import { handleMetadataPacket } from "~/utils/packets";
 import { startSpan } from "~/v3/tracing.server";
-import {
+import type {
   TriggerTaskServiceOptions,
   TriggerTaskServiceResult,
 } from "../../v3/services/triggerTask.server";

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

@@ -88,6 +88,8 @@ export class TriggerTaskService extends WithRunEngine {
     body: TriggerTaskRequestBody,
     options: TriggerTaskServiceOptions = {}
   ): Promise<TriggerTaskServiceResult | undefined> {
+    const traceEventConcern = new DefaultTraceEventsConcern(eventRepository);
+
     const service = new RunEngineTriggerTaskService({
       prisma: this._prisma,
       engine: this._engine,
@@ -97,10 +99,10 @@ export class TriggerTaskService extends WithRunEngine {
       idempotencyKeyConcern: new IdempotencyKeyConcern(
         this._prisma,
         this._engine,
-        new DefaultTraceEventsConcern(eventRepository)
+        traceEventConcern
       ),
       runNumberIncrementer: new DefaultRunNumberIncrementer(),
-      traceEventConcern: new DefaultTraceEventsConcern(eventRepository),
+      traceEventConcern,
       runChainStateManager: new DefaultRunChainStateManager(
         this._prisma,
         env.RUN_ENGINE_RELEASE_CONCURRENCY_ENABLED === "1"

apps/webapp/test/engine/triggerTask.test.ts

@@ -1,11 +1,23 @@
+import { describe, expect, vi } from "vitest";
+
+// Mock the db prisma client
+vi.mock("~/db.server", () => ({
+  prisma: {},
+}));
+
+vi.mock("~/services/platform.v3.server", () => ({
+  getEntitlement: vi.fn(),
+}));
+
 import { RunEngine } from "@internal/run-engine";
 import { setupAuthenticatedEnvironment, setupBackgroundWorker } from "@internal/run-engine/tests";
 import { containerTest } from "@internal/testcontainers";
 import { trace } from "@opentelemetry/api";
 import { IOPacket } from "@trigger.dev/core/v3";
-import { describe, expect, vi } from "vitest";
+import { TaskRun } from "@trigger.dev/database";
 import { IdempotencyKeyConcern } from "~/runEngine/concerns/idempotencyKeys.server";
 import { DefaultQueueManager } from "~/runEngine/concerns/queues.server";
+import { DefaultRunChainStateManager } from "~/runEngine/concerns/runChainStates.server";
 import {
   EntitlementValidationParams,
   MaxAttemptsValidationParams,
@@ -20,8 +32,7 @@ import {
   ValidationResult,
 } from "~/runEngine/types";
 import { RunEngineTriggerTaskService } from "../../app/runEngine/services/triggerTask.server";
-import { TaskRun, TaskRunStatus } from "@trigger.dev/database";
-import { DefaultRunChainStateManager } from "~/runEngine/concerns/runChainStates.server";
+import { getEntitlement } from "~/services/platform.v3.server";
 
 vi.setConfig({ testTimeout: 30_000 }); // 30 seconds timeout