add snapshots since methods and route

Author: nicktrn

apps/webapp/app/routes/engine.v1.worker-actions.runs.$runFriendlyId.snapshots.since.$snapshotId.ts

@@ -0,0 +1,30 @@
+import { json, TypedResponse } from "@remix-run/server-runtime";
+import { WorkerApiRunSnapshotsSinceResponseBody } from "@trigger.dev/core/v3/workers";
+import { z } from "zod";
+import { createLoaderWorkerApiRoute } from "~/services/routeBuilders/apiBuilder.server";
+
+export const loader = createLoaderWorkerApiRoute(
+  {
+    params: z.object({
+      runFriendlyId: z.string(),
+      snapshotId: z.string(),
+    }),
+  },
+  async ({
+    authenticatedWorker,
+    params,
+  }): Promise<TypedResponse<WorkerApiRunSnapshotsSinceResponseBody>> => {
+    const { runFriendlyId, snapshotId } = params;
+
+    const executions = await authenticatedWorker.getSnapshotsSince({
+      runFriendlyId,
+      snapshotId,
+    });
+
+    if (!executions) {
+      throw new Error("Failed to retrieve snapshots since given snapshot");
+    }
+
+    return json({ executions });
+  }
+);

apps/webapp/app/v3/services/worker/workerGroupTokenService.server.ts

@@ -759,6 +759,19 @@ export class AuthenticatedWorkerInstance extends WithRunEngine {
     });
   }
 
+  async getSnapshotsSince({
+    runFriendlyId,
+    snapshotId,
+  }: {
+    runFriendlyId: string;
+    snapshotId: string;
+  }) {
+    return await this._engine.getSnapshotsSince({
+      runId: fromFriendlyId(runFriendlyId),
+      snapshotId: fromFriendlyId(snapshotId),
+    });
+  }
+
   toJSON(): WorkerGroupTokenAuthenticationResponse {
     if (this.type === WorkerInstanceGroupType.MANAGED) {
       return {

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

@@ -43,6 +43,8 @@ import { EnqueueSystem } from "./systems/enqueueSystem.js";
 import {
   ExecutionSnapshotSystem,
   getLatestExecutionSnapshot,
+  getExecutionSnapshotsSince,
+  executionDataFromSnapshot,
 } from "./systems/executionSnapshotSystem.js";
 import { PendingVersionSystem } from "./systems/pendingVersionSystem.js";
 import { ReleaseConcurrencySystem } from "./systems/releaseConcurrencySystem.js";
@@ -1100,43 +1102,31 @@ export class RunEngine {
     const prisma = tx ?? this.prisma;
     try {
       const snapshot = await getLatestExecutionSnapshot(prisma, runId);
+      return executionDataFromSnapshot(snapshot);
+    } catch (e) {
+      this.logger.error("Failed to getRunExecutionData", {
+        message: e instanceof Error ? e.message : e,
+      });
+      return null;
+    }
+  }
 
-      const executionData: RunExecutionData = {
-        version: "1" as const,
-        snapshot: {
-          id: snapshot.id,
-          friendlyId: snapshot.friendlyId,
-          executionStatus: snapshot.executionStatus,
-          description: snapshot.description,
-        },
-        run: {
-          id: snapshot.runId,
-          friendlyId: snapshot.runFriendlyId,
-          status: snapshot.runStatus,
-          attemptNumber: snapshot.attemptNumber ?? undefined,
-        },
-        batch: snapshot.batchId
-          ? {
-              id: snapshot.batchId,
-              friendlyId: BatchId.toFriendlyId(snapshot.batchId),
-            }
-          : undefined,
-        checkpoint: snapshot.checkpoint
-          ? {
-              id: snapshot.checkpoint.id,
-              friendlyId: snapshot.checkpoint.friendlyId,
-              type: snapshot.checkpoint.type,
-              location: snapshot.checkpoint.location,
-              imageRef: snapshot.checkpoint.imageRef,
-              reason: snapshot.checkpoint.reason ?? undefined,
-            }
-          : undefined,
-        completedWaitpoints: snapshot.completedWaitpoints,
-      };
+  async getSnapshotsSince({
+    runId,
+    snapshotId,
+    tx,
+  }: {
+    runId: string;
+    snapshotId: string;
+    tx?: PrismaClientOrTransaction;
+  }): Promise<RunExecutionData[] | null> {
+    const prisma = tx ?? this.prisma;
 
-      return executionData;
+    try {
+      const snapshots = await getExecutionSnapshotsSince(prisma, runId, snapshotId);
+      return snapshots.map(executionDataFromSnapshot);
     } catch (e) {
-      this.logger.error("Failed to getRunExecutionData", {
+      this.logger.error("Failed to getSnapshotsSince", {
         message: e instanceof Error ? e.message : e,
       });
       return null;
@@ -1158,9 +1148,6 @@ export class RunEngine {
     }
   }
 
-  //#endregion
-
-  //#region Heartbeat
   async #handleStalledSnapshot({
     runId,
     snapshotId,

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

@@ -1,4 +1,4 @@
-import { CompletedWaitpoint, ExecutionResult } from "@trigger.dev/core/v3";
+import { CompletedWaitpoint, ExecutionResult, RunExecutionData } from "@trigger.dev/core/v3";
 import { BatchId, RunId, SnapshotId } from "@trigger.dev/core/v3/isomorphic";
 import {
   Prisma,
@@ -17,31 +17,23 @@ export type ExecutionSnapshotSystemOptions = {
   heartbeatTimeouts: HeartbeatTimeouts;
 };
 
-export interface LatestExecutionSnapshot extends TaskRunExecutionSnapshot {
+export interface EnhancedExecutionSnapshot extends TaskRunExecutionSnapshot {
   friendlyId: string;
   runFriendlyId: string;
   checkpoint: TaskRunCheckpoint | null;
   completedWaitpoints: CompletedWaitpoint[];
 }
 
-/* Gets the most recent valid snapshot for a run */
-export async function getLatestExecutionSnapshot(
-  prisma: PrismaClientOrTransaction,
-  runId: string
-): Promise<LatestExecutionSnapshot> {
-  const snapshot = await prisma.taskRunExecutionSnapshot.findFirst({
-    where: { runId, isValid: true },
-    include: {
-      completedWaitpoints: true,
-      checkpoint: true,
-    },
-    orderBy: { createdAt: "desc" },
-  });
-
-  if (!snapshot) {
-    throw new Error(`No execution snapshot found for TaskRun ${runId}`);
-  }
+type ExecutionSnapshotWithCheckAndWaitpoints = Prisma.TaskRunExecutionSnapshotGetPayload<{
+  include: {
+    checkpoint: true;
+    completedWaitpoints: true;
+  };
+}>;
 
+function enhanceExecutionSnapshot(
+  snapshot: ExecutionSnapshotWithCheckAndWaitpoints
+): EnhancedExecutionSnapshot {
   return {
     ...snapshot,
     friendlyId: SnapshotId.toFriendlyId(snapshot.id),
@@ -99,6 +91,27 @@ export async function getLatestExecutionSnapshot(
   };
 }
 
+/* Gets the most recent valid snapshot for a run */
+export async function getLatestExecutionSnapshot(
+  prisma: PrismaClientOrTransaction,
+  runId: string
+): Promise<EnhancedExecutionSnapshot> {
+  const snapshot = await prisma.taskRunExecutionSnapshot.findFirst({
+    where: { runId, isValid: true },
+    include: {
+      completedWaitpoints: true,
+      checkpoint: true,
+    },
+    orderBy: { createdAt: "desc" },
+  });
+
+  if (!snapshot) {
+    throw new Error(`No execution snapshot found for TaskRun ${runId}`);
+  }
+
+  return enhanceExecutionSnapshot(snapshot);
+}
+
 export async function getExecutionSnapshotCompletedWaitpoints(
   prisma: PrismaClientOrTransaction,
   snapshotId: string
@@ -141,6 +154,72 @@ export function executionResultFromSnapshot(snapshot: TaskRunExecutionSnapshot):
   };
 }
 
+export function executionDataFromSnapshot(snapshot: EnhancedExecutionSnapshot): RunExecutionData {
+  return {
+    version: "1" as const,
+    snapshot: {
+      id: snapshot.id,
+      friendlyId: snapshot.friendlyId,
+      executionStatus: snapshot.executionStatus,
+      description: snapshot.description,
+    },
+    run: {
+      id: snapshot.runId,
+      friendlyId: snapshot.runFriendlyId,
+      status: snapshot.runStatus,
+      attemptNumber: snapshot.attemptNumber ?? undefined,
+    },
+    batch: snapshot.batchId
+      ? {
+          id: snapshot.batchId,
+          friendlyId: BatchId.toFriendlyId(snapshot.batchId),
+        }
+      : undefined,
+    checkpoint: snapshot.checkpoint
+      ? {
+          id: snapshot.checkpoint.id,
+          friendlyId: snapshot.checkpoint.friendlyId,
+          type: snapshot.checkpoint.type,
+          location: snapshot.checkpoint.location,
+          imageRef: snapshot.checkpoint.imageRef,
+          reason: snapshot.checkpoint.reason ?? undefined,
+        }
+      : undefined,
+    completedWaitpoints: snapshot.completedWaitpoints,
+  };
+}
+
+export async function getExecutionSnapshotsSince(
+  prisma: PrismaClientOrTransaction,
+  runId: string,
+  sinceSnapshotId: string
+): Promise<EnhancedExecutionSnapshot[]> {
+  // Find the createdAt of the sinceSnapshotId
+  const sinceSnapshot = await prisma.taskRunExecutionSnapshot.findUnique({
+    where: { id: sinceSnapshotId },
+    select: { createdAt: true },
+  });
+
+  if (!sinceSnapshot) {
+    throw new Error(`No execution snapshot found for id ${sinceSnapshotId}`);
+  }
+
+  const snapshots = await prisma.taskRunExecutionSnapshot.findMany({
+    where: {
+      runId,
+      isValid: true,
+      createdAt: { gt: sinceSnapshot.createdAt },
+    },
+    include: {
+      completedWaitpoints: true,
+      checkpoint: true,
+    },
+    orderBy: { createdAt: "asc" },
+  });
+
+  return snapshots.map(enhanceExecutionSnapshot);
+}
+
 export class ExecutionSnapshotSystem {
   private readonly $: SystemResources;
   private readonly heartbeatTimeouts: HeartbeatTimeouts;