add snapshot poller service

Author: nicktrn

packages/cli-v3/src/entryPoints/managed-run-controller.ts

@@ -3,7 +3,6 @@ import { env as stdEnv } from "std-env";
 import { readJSONFile } from "../utilities/fileSystem.js";
 import {
   type CompleteRunAttemptResult,
-  HeartbeatService,
   type RunExecutionData,
   SuspendedProcessError,
   type TaskRunExecutionMetrics,
@@ -26,6 +25,7 @@ import { RunnerEnv } from "./managed/env.js";
 import { MetadataClient } from "./managed/overrides.js";
 import { RunLogger, SendDebugLogOptions } from "./managed/logger.js";
 import { RunExecutionHeartbeat } from "./managed/heartbeat.js";
+import { RunExecutionSnapshotPoller } from "./managed/poller.js";
 
 const env = new RunnerEnv(stdEnv);
 
@@ -54,8 +54,8 @@ class ManagedRunController {
   private socket: Socket<WorkloadServerToClientEvents, WorkloadClientToServerEvents>;
   private readonly logger: RunLogger;
 
-  private readonly runHeartbeat: HeartbeatService;
-  private readonly snapshotPoller: RunExecutionHeartbeat;
+  private readonly runHeartbeat: RunExecutionHeartbeat;
+  private readonly snapshotPoller: RunExecutionSnapshotPoller;
 
   constructor(opts: ManagedRunControllerOptions) {
     this.workerManifest = opts.workerManifest;
@@ -100,55 +100,25 @@ class ManagedRunController {
       });
     }
 
-    this.snapshotPoller = new RunExecutionHeartbeat({
+    this.snapshotPoller = new RunExecutionSnapshotPoller({
       // @ts-expect-error
       runFriendlyId: env.TRIGGER_RUN_ID,
       // @ts-expect-error
       snapshotFriendlyId: env.TRIGGER_SNAPSHOT_ID,
       httpClient: this.httpClient,
       logger: this.logger,
-      heartbeatIntervalSeconds: this.heartbeatIntervalSeconds,
+      snapshotPollIntervalSeconds: this.snapshotPollIntervalSeconds,
+      handleSnapshotChange: this.handleSnapshotChange,
     });
 
-    this.runHeartbeat = new HeartbeatService({
-      heartbeat: async () => {
-        if (!this.runFriendlyId || !this.snapshotFriendlyId) {
-          this.sendDebugLog({
-            runId: this.runFriendlyId,
-            message: "Skipping heartbeat, no run ID or snapshot ID",
-          });
-          return;
-        }
-
-        this.sendDebugLog({
-          runId: this.runFriendlyId,
-          message: "heartbeat: started",
-        });
-
-        const response = await this.httpClient.heartbeatRun(
-          this.runFriendlyId,
-          this.snapshotFriendlyId
-        );
-
-        if (!response.success) {
-          this.sendDebugLog({
-            runId: this.runFriendlyId,
-            message: "heartbeat: failed",
-            properties: {
-              error: response.error,
-            },
-          });
-        }
-      },
-      intervalMs: this.heartbeatIntervalSeconds * 1000,
-      leadingEdge: false,
-      onError: async (error) => {
-        this.sendDebugLog({
-          runId: this.runFriendlyId,
-          message: "Failed to send heartbeat",
-          properties: { error: error instanceof Error ? error.message : String(error) },
-        });
-      },
+    this.runHeartbeat = new RunExecutionHeartbeat({
+      // @ts-expect-error
+      runFriendlyId: env.TRIGGER_RUN_ID,
+      // @ts-expect-error
+      snapshotFriendlyId: env.TRIGGER_SNAPSHOT_ID,
+      httpClient: this.httpClient,
+      logger: this.logger,
+      heartbeatIntervalSeconds: this.heartbeatIntervalSeconds,
     });
 
     process.on("SIGTERM", async () => {
@@ -429,6 +399,8 @@ class ManagedRunController {
 
       try {
         this.updateRunPhase(run, snapshot);
+
+        this.runHeartbeat.updateSnapshotId(snapshot.friendlyId);
         this.snapshotPoller.updateSnapshotId(snapshot.friendlyId);
       } catch (error) {
         this.sendDebugLog({

packages/cli-v3/src/entryPoints/managed/poller.ts

@@ -0,0 +1,114 @@
+import { WorkloadHttpClient } from "@trigger.dev/core/v3/runEngineWorker";
+import { RunLogger } from "./logger.js";
+import { HeartbeatService, RunExecutionData } from "@trigger.dev/core/v3";
+
+export type RunExecutionSnapshotPollerOptions = {
+  runFriendlyId: string;
+  snapshotFriendlyId: string;
+  httpClient: WorkloadHttpClient;
+  logger: RunLogger;
+  snapshotPollIntervalSeconds: number;
+  handleSnapshotChange: (execution: RunExecutionData) => Promise<void>;
+};
+
+export class RunExecutionSnapshotPoller {
+  private readonly logger: RunLogger;
+  private readonly poller: HeartbeatService;
+  private readonly httpClient: WorkloadHttpClient;
+
+  private readonly runFriendlyId: string;
+  private snapshotFriendlyId: string;
+
+  private readonly snapshotPollIntervalSeconds: number;
+
+  private readonly handleSnapshotChange: (execution: RunExecutionData) => Promise<void>;
+
+  constructor(opts: RunExecutionSnapshotPollerOptions) {
+    this.logger = opts.logger;
+    this.httpClient = opts.httpClient;
+
+    this.runFriendlyId = opts.runFriendlyId;
+    this.snapshotFriendlyId = opts.snapshotFriendlyId;
+
+    this.handleSnapshotChange = opts.handleSnapshotChange;
+
+    this.poller = new HeartbeatService({
+      heartbeat: async () => {
+        if (!this.runFriendlyId) {
+          this.logger.sendDebugLog({
+            runId: this.runFriendlyId,
+            message: "Skipping snapshot poll, no run ID",
+          });
+          return;
+        }
+
+        this.logger.sendDebugLog({
+          runId: this.runFriendlyId,
+          message: "Polling for latest snapshot",
+        });
+
+        this.logger.sendDebugLog({
+          runId: this.runFriendlyId,
+          message: `snapshot poll: started`,
+          properties: {
+            snapshotId: this.snapshotFriendlyId,
+          },
+        });
+
+        const response = await this.httpClient.getRunExecutionData(this.runFriendlyId);
+
+        if (!response.success) {
+          this.logger.sendDebugLog({
+            runId: this.runFriendlyId,
+            message: "Snapshot poll failed",
+            properties: {
+              error: response.error,
+            },
+          });
+
+          this.logger.sendDebugLog({
+            runId: this.runFriendlyId,
+            message: `snapshot poll: failed`,
+            properties: {
+              snapshotId: this.snapshotFriendlyId,
+              error: response.error,
+            },
+          });
+
+          return;
+        }
+
+        await this.handleSnapshotChange(response.data.execution);
+      },
+      intervalMs: this.snapshotPollIntervalSeconds * 1000,
+      leadingEdge: false,
+      onError: async (error) => {
+        this.logger.sendDebugLog({
+          runId: this.runFriendlyId,
+          message: "Failed to poll for snapshot",
+          properties: { error: error instanceof Error ? error.message : String(error) },
+        });
+      },
+    });
+  }
+
+  resetCurrentInterval() {
+    this.poller.resetCurrentInterval();
+  }
+
+  updateSnapshotId(snapshotId: string) {
+    this.snapshotFriendlyId = snapshotId;
+  }
+
+  updateInterval(intervalMs: number) {
+    this.poller.updateInterval(intervalMs);
+  }
+
+  start() {
+    this.poller.start();
+  }
+
+  stop() {
+    this.poller.stop();
+  }
+}