v3: checkpoint and reliability improvements (#1198)

* only checkpoint retries with delays greater than threshold

* rename checkpoint threshold env var

* log task monitor ignores

* crash runs with unbounded attempts

* fix retry check in shared queue consumer

* add missing stop for env var sync spinner

* prod entry point refactor

* missing awaits

* more verbose prod flush and exit logs

* reduce checkpoint support logs

* heartbeat while checkpointing between retries

* dynamic coordinator config

* measure lazy attempt creation time in prod

* simplify delay threshold

* heartbeat clarifications

* crash run if it doesn't reach checkpointable state

* require dynamic config threshold

* fix retry prep, await previous worker kill

* unify wait mechanics

* fix prod worker without tasks error

* ensure worker is ready to be checkpointed for dependency waits

* improve worker attempt creation logging

* prevent crashes caused by failed socket schema parsing

* fix dynamic imports in v3 catalog

* clarify attempt retry mechanics

* move backoff helper to core-apps

* remove core-apps barrel file

* add backoff execute with callback

* deprecate non-lazy attempt messages

* update socket.io-client to v4.7.5

* fix socket.io types for emits with timeout

* retry all the things

* remove todo

* fix retry restores

* improve index failure logs

* retry incomplete dependency waits

* fix checkpoint in-progress detection

* prevent losing messages during reconnect

* checkpoint when greater or equal to threshold

* improve handling of duration wait edge cases

* add ready for lazy attempt replay

* retry attempt completion

* allow failing runs with unfriendly run id

* fix min max jitter

* cancel checkpoints on run failure

* improve attempt creation errors

* prevent crashing run on failed cleanup

* handle at-least-once execute lazy attempt delivery

* log exit code on prepare for retry

* fix timeout promise

* mark some things

* chaos monkey superpowers

* refactor checkpointer

* set chaos monkey defaults

* less chaos

* fix backoff

* handle uncaught entry point exceptions

* only replay rpcs on true reconnects

* allow resume unless final run status

* add changeset

* small fixes

.changeset/mighty-eggs-grab.md

@@ -0,0 +1,30 @@
+---
+"@trigger.dev/core-apps": patch
+"trigger.dev": patch
+"@trigger.dev/core": patch
+---
+
+Tasks should now be much more robust and resilient to reconnects during crucial operations and other failure scenarios.
+
+Task runs now have to signal checkpointable state prior to ALL checkpoints. This ensures flushing always happens.
+
+All important socket.io RPCs will now be retried with backoff. Actions relying on checkpoints will be replayed if we haven't been checkpointed and restored as expected, e.g. after reconnect.
+
+Other changes:
+
+- Fix retry check in shared queue
+- Fix env var sync spinner
+- Heartbeat between retries
+- Fix retry prep
+- Fix prod worker no tasks detection
+- Fail runs above `MAX_TASK_RUN_ATTEMPTS`
+- Additional debug logs in all places
+- Prevent crashes due to failed socket schema parsing
+- Remove core-apps barrel
+- Upgrade socket.io-client to fix an ACK memleak
+- Additional index failure logs
+- Prevent message loss during reconnect
+- Prevent burst of heartbeats on reconnect
+- Prevent crash on failed cleanup
+- Handle at-least-once lazy execute message delivery
+- Handle uncaught entry point exceptions

.env.example

@@ -71,7 +71,7 @@ COORDINATOR_SECRET=coordinator-secret # generate the actual secret with `openssl
 # OBJECT_STORE_BASE_URL="https://{bucket}.{accountId}.r2.cloudflarestorage.com"
 # OBJECT_STORE_ACCESS_KEY_ID=
 # OBJECT_STORE_SECRET_ACCESS_KEY=
-# RUNTIME_WAIT_THRESHOLD_IN_MS=10000
+# CHECKPOINT_THRESHOLD_IN_MS=10000
 
 # These control the server-side internal telemetry
 # INTERNAL_OTEL_TRACE_EXPORTER_URL=<URL to send traces to>

apps/coordinator/package.json

@@ -21,8 +21,7 @@
     "execa": "^8.0.1",
     "nanoid": "^5.0.6",
     "prom-client": "^15.1.0",
-    "socket.io": "4.7.4",
-    "socket.io-client": "4.7.4"
+    "socket.io": "4.7.4"
   },
   "devDependencies": {
     "@types/node": "^18",

apps/coordinator/src/chaosMonkey.ts

@@ -0,0 +1,95 @@
+import type { Execa$ } from "execa";
+import { setTimeout as timeout } from "node:timers/promises";
+
+class ChaosMonkeyError extends Error {
+  constructor(message: string) {
+    super(message);
+    this.name = "ChaosMonkeyError";
+  }
+}
+
+export class ChaosMonkey {
+  private chaosEventRate = 0.2;
+  private delayInSeconds = 45;
+
+  constructor(private enabled = false) {
+    if (this.enabled) {
+      console.log("🍌 Chaos monkey enabled");
+    }
+  }
+
+  static Error = ChaosMonkeyError;
+
+  enable() {
+    this.enabled = true;
+    console.log("🍌 Chaos monkey enabled");
+  }
+
+  disable() {
+    this.enabled = false;
+    console.log("🍌 Chaos monkey disabled");
+  }
+
+  async call({
+    $,
+    throwErrors = true,
+    addDelays = true,
+  }: {
+    $?: Execa$<string>;
+    throwErrors?: boolean;
+    addDelays?: boolean;
+  } = {}) {
+    if (!this.enabled) {
+      return;
+    }
+
+    const random = Math.random();
+
+    if (random > this.chaosEventRate) {
+      // Don't interfere with normal operation
+      return;
+    }
+
+    const chaosEvents: Array<() => Promise<any>> = [];
+
+    if (addDelays) {
+      chaosEvents.push(async () => {
+        console.log("🍌 Chaos monkey: Add delay");
+
+        if ($) {
+          await $`sleep ${this.delayInSeconds}`;
+        } else {
+          await timeout(this.delayInSeconds * 1000);
+        }
+      });
+    }
+
+    if (throwErrors) {
+      chaosEvents.push(async () => {
+        console.log("🍌 Chaos monkey: Throw error");
+
+        if ($) {
+          await $`false`;
+        } else {
+          throw new ChaosMonkey.Error("🍌 Chaos monkey: Throw error");
+        }
+      });
+    }
+
+    if (chaosEvents.length === 0) {
+      console.error("🍌 Chaos monkey: No events selected");
+      return;
+    }
+
+    const randomIndex = Math.floor(Math.random() * chaosEvents.length);
+
+    const chaosEvent = chaosEvents[randomIndex];
+
+    if (!chaosEvent) {
+      console.error("🍌 Chaos monkey: No event found");
+      return;
+    }
+
+    await chaosEvent();
+  }
+}

apps/coordinator/src/checkpointer.ts

@@ -0,0 +1,587 @@
+import { ExponentialBackoff } from "@trigger.dev/core-apps/backoff";
+import { isExecaChildProcess, testDockerCheckpoint } from "@trigger.dev/core-apps/checkpoints";
+import { SimpleLogger } from "@trigger.dev/core-apps/logger";
+import { $ } from "execa";
+import { nanoid } from "nanoid";
+import fs from "node:fs/promises";
+import { ChaosMonkey } from "./chaosMonkey";
+
+type CheckpointerInitializeReturn = {
+  canCheckpoint: boolean;
+  willSimulate: boolean;
+};
+
+type CheckpointAndPushOptions = {
+  runId: string;
+  leaveRunning?: boolean;
+  projectRef: string;
+  deploymentVersion: string;
+  shouldHeartbeat?: boolean;
+};
+
+type CheckpointAndPushResult =
+  | { success: true; checkpoint: CheckpointData }
+  | {
+      success: false;
+      reason?: "CANCELED" | "DISABLED" | "ERROR" | "IN_PROGRESS" | "NO_SUPPORT" | "SKIP_RETRYING";
+    };
+
+type CheckpointData = {
+  location: string;
+  docker: boolean;
+};
+
+type CheckpointerOptions = {
+  dockerMode: boolean;
+  forceSimulate: boolean;
+  heartbeat: (runId: string) => void;
+  registryHost?: string;
+  registryNamespace?: string;
+  registryTlsVerify?: boolean;
+  disableCheckpointSupport?: boolean;
+  checkpointPath?: string;
+  simulateCheckpointFailure?: boolean;
+  simulateCheckpointFailureSeconds?: number;
+  simulatePushFailure?: boolean;
+  simulatePushFailureSeconds?: number;
+  chaosMonkey?: ChaosMonkey;
+};
+
+async function getFileSize(filePath: string): Promise<number> {
+  try {
+    const stats = await fs.stat(filePath);
+    return stats.size;
+  } catch (error) {
+    console.error("Error getting file size:", error);
+    return -1;
+  }
+}
+
+async function getParsedFileSize(filePath: string) {
+  const sizeInBytes = await getFileSize(filePath);
+
+  let message = `Size in bytes: ${sizeInBytes}`;
+
+  if (sizeInBytes > 1024 * 1024) {
+    const sizeInMB = (sizeInBytes / 1024 / 1024).toFixed(2);
+    message = `Size in MB (rounded): ${sizeInMB}`;
+  } else if (sizeInBytes > 1024) {
+    const sizeInKB = (sizeInBytes / 1024).toFixed(2);
+    message = `Size in KB (rounded): ${sizeInKB}`;
+  }
+
+  return {
+    path: filePath,
+    sizeInBytes,
+    message,
+  };
+}
+
+export class Checkpointer {
+  #initialized = false;
+  #canCheckpoint = false;
+  #dockerMode: boolean;
+
+  #logger = new SimpleLogger("[checkptr]");
+  #abortControllers = new Map<string, AbortController>();
+  #failedCheckpoints = new Map<string, unknown>();
+  #waitingForRetry = new Set<string>();
+
+  private registryHost: string;
+  private registryNamespace: string;
+  private registryTlsVerify: boolean;
+
+  private disableCheckpointSupport: boolean;
+  private checkpointPath: string;
+
+  private simulateCheckpointFailure: boolean;
+  private simulateCheckpointFailureSeconds: number;
+  private simulatePushFailure: boolean;
+  private simulatePushFailureSeconds: number;
+
+  private chaosMonkey: ChaosMonkey;
+
+  constructor(private opts: CheckpointerOptions) {
+    this.#dockerMode = opts.dockerMode;
+
+    this.registryHost = opts.registryHost ?? "localhost:5000";
+    this.registryNamespace = opts.registryNamespace ?? "trigger";
+    this.registryTlsVerify = opts.registryTlsVerify ?? true;
+
+    this.disableCheckpointSupport = opts.disableCheckpointSupport ?? false;
+    this.checkpointPath = opts.checkpointPath ?? "/checkpoints";
+
+    this.simulateCheckpointFailure = opts.simulateCheckpointFailure ?? false;
+    this.simulateCheckpointFailureSeconds = opts.simulateCheckpointFailureSeconds ?? 300;
+    this.simulatePushFailure = opts.simulatePushFailure ?? false;
+    this.simulatePushFailureSeconds = opts.simulatePushFailureSeconds ?? 300;
+
+    this.chaosMonkey = opts.chaosMonkey ?? new ChaosMonkey(!!process.env.CHAOS_MONKEY_ENABLED);
+  }
+
+  async init(): Promise<CheckpointerInitializeReturn> {
+    if (this.#initialized) {
+      return this.#getInitReturn(this.#canCheckpoint);
+    }
+
+    this.#logger.log(`${this.#dockerMode ? "Docker" : "Kubernetes"} mode`);
+
+    if (this.#dockerMode) {
+      const testCheckpoint = await testDockerCheckpoint();
+
+      if (testCheckpoint.ok) {
+        return this.#getInitReturn(true);
+      }
+
+      this.#logger.error(testCheckpoint.message, testCheckpoint.error ?? "");
+      return this.#getInitReturn(false);
+    } else {
+      try {
+        await $`buildah login --get-login ${this.registryHost}`;
+      } catch (error) {
+        this.#logger.error(`No checkpoint support: Not logged in to registry ${this.registryHost}`);
+        return this.#getInitReturn(false);
+      }
+    }
+
+    return this.#getInitReturn(true);
+  }
+
+  #getInitReturn(canCheckpoint: boolean): CheckpointerInitializeReturn {
+    this.#canCheckpoint = canCheckpoint;
+
+    if (canCheckpoint) {
+      if (!this.#initialized) {
+        this.#logger.log("Full checkpoint support!");
+      }
+    }
+
+    this.#initialized = true;
+
+    const willSimulate = this.#dockerMode && (!this.#canCheckpoint || this.opts.forceSimulate);
+
+    if (willSimulate) {
+      this.#logger.log("Simulation mode enabled. Containers will be paused, not checkpointed.", {
+        forceSimulate: this.opts.forceSimulate,
+      });
+    }
+
+    return {
+      canCheckpoint,
+      willSimulate,
+    };
+  }
+
+  #getImageRef(projectRef: string, deploymentVersion: string, shortCode: string) {
+    return `${this.registryHost}/${this.registryNamespace}/${projectRef}:${deploymentVersion}.prod-${shortCode}`;
+  }
+
+  #getExportLocation(projectRef: string, deploymentVersion: string, shortCode: string) {
+    const basename = `${projectRef}-${deploymentVersion}-${shortCode}`;
+
+    if (this.#dockerMode) {
+      return basename;
+    } else {
+      return `${this.checkpointPath}/${basename}.tar`;
+    }
+  }
+
+  async checkpointAndPush(opts: CheckpointAndPushOptions): Promise<CheckpointData | undefined> {
+    const start = performance.now();
+    this.#logger.log(`checkpointAndPush() start`, { start, opts });
+
+    let interval: NodeJS.Timer | undefined;
+
+    if (opts.shouldHeartbeat) {
+      interval = setInterval(() => {
+        this.#logger.log("Sending heartbeat", { runId: opts.runId });
+        this.opts.heartbeat(opts.runId);
+      }, 20_000);
+    }
+
+    try {
+      const result = await this.#checkpointAndPushWithBackoff(opts);
+
+      const end = performance.now();
+      this.#logger.log(`checkpointAndPush() end`, {
+        start,
+        end,
+        diff: end - start,
+        opts,
+        success: result.success,
+      });
+
+      if (!result.success) {
+        return;
+      }
+
+      return result.checkpoint;
+    } finally {
+      if (opts.shouldHeartbeat) {
+        clearInterval(interval);
+      }
+    }
+  }
+
+  isCheckpointing(runId: string) {
+    return this.#abortControllers.has(runId) || this.#waitingForRetry.has(runId);
+  }
+
+  cancelCheckpoint(runId: string): boolean {
+    // If the last checkpoint failed, pretend we canceled it
+    // This ensures tasks don't wait for external resume messages to continue
+    if (this.#hasFailedCheckpoint(runId)) {
+      this.#clearFailedCheckpoint(runId);
+      return true;
+    }
+
+    if (this.#waitingForRetry.has(runId)) {
+      this.#waitingForRetry.delete(runId);
+      return true;
+    }
+
+    const controller = this.#abortControllers.get(runId);
+
+    if (!controller) {
+      this.#logger.debug("Nothing to cancel", { runId });
+      return false;
+    }
+
+    controller.abort("cancelCheckpointing()");
+    this.#abortControllers.delete(runId);
+
+    return true;
+  }
+
+  async #checkpointAndPushWithBackoff({
+    runId,
+    leaveRunning = true, // This mirrors kubernetes behaviour more accurately
+    projectRef,
+    deploymentVersion,
+  }: CheckpointAndPushOptions): Promise<CheckpointAndPushResult> {
+    this.#logger.log("Checkpointing with backoff", {
+      runId,
+      leaveRunning,
+      projectRef,
+      deploymentVersion,
+    });
+
+    const backoff = new ExponentialBackoff()
+      .type("EqualJitter")
+      .base(3)
+      .max(3 * 3600)
+      .maxElapsed(48 * 3600);
+
+    for await (const { delay, retry } of backoff) {
+      try {
+        if (retry > 0) {
+          this.#logger.error("Retrying checkpoint", {
+            runId,
+            retry,
+            delay,
+          });
+
+          this.#waitingForRetry.add(runId);
+          await new Promise((resolve) => setTimeout(resolve, delay.milliseconds));
+
+          if (!this.#waitingForRetry.has(runId)) {
+            this.#logger.log("Checkpoint canceled while waiting for retry", { runId });
+            return { success: false, reason: "CANCELED" };
+          } else {
+            this.#waitingForRetry.delete(runId);
+          }
+        }
+
+        const result = await this.#checkpointAndPush({
+          runId,
+          leaveRunning,
+          projectRef,
+          deploymentVersion,
+        });
+
+        if (result.success) {
+          return result;
+        }
+
+        if (result.reason === "CANCELED") {
+          this.#logger.log("Checkpoint canceled, won't retry", { runId });
+          // Don't fail the checkpoint, as it was canceled
+          return result;
+        }
+
+        if (result.reason === "IN_PROGRESS") {
+          this.#logger.log("Checkpoint already in progress, won't retry", { runId });
+          this.#failCheckpoint(runId, result.reason);
+          return result;
+        }
+
+        if (result.reason === "NO_SUPPORT") {
+          this.#logger.log("No checkpoint support, won't retry", { runId });
+          this.#failCheckpoint(runId, result.reason);
+          return result;
+        }
+
+        if (result.reason === "DISABLED") {
+          this.#logger.log("Checkpoint support disabled, won't retry", { runId });
+          this.#failCheckpoint(runId, result.reason);
+          return result;
+        }
+
+        if (result.reason === "SKIP_RETRYING") {
+          this.#logger.log("Skipping retrying", { runId });
+          return result;
+        }
+
+        continue;
+      } catch (error) {
+        this.#logger.error("Checkpoint error", {
+          retry,
+          runId,
+          delay,
+          error: error instanceof Error ? error.message : error,
+        });
+      }
+    }
+
+    this.#logger.error(`Checkpoint failed after exponential backoff`, {
+      runId,
+      leaveRunning,
+      projectRef,
+      deploymentVersion,
+    });
+    this.#failCheckpoint(runId, "ERROR");
+
+    return { success: false, reason: "ERROR" };
+  }
+
+  async #checkpointAndPush({
+    runId,
+    leaveRunning = true, // This mirrors kubernetes behaviour more accurately
+    projectRef,
+    deploymentVersion,
+  }: CheckpointAndPushOptions): Promise<CheckpointAndPushResult> {
+    await this.init();
+
+    const options = {
+      runId,
+      leaveRunning,
+      projectRef,
+      deploymentVersion,
+    };
+
+    if (!this.#dockerMode && !this.#canCheckpoint) {
+      this.#logger.error("No checkpoint support. Simulation requires docker.");
+      return { success: false, reason: "NO_SUPPORT" };
+    }
+
+    if (this.isCheckpointing(runId)) {
+      this.#logger.error("Checkpoint procedure already in progress", { options });
+      return { success: false, reason: "IN_PROGRESS" };
+    }
+
+    // This is a new checkpoint, clear any last failure for this run
+    this.#clearFailedCheckpoint(runId);
+
+    if (this.disableCheckpointSupport) {
+      this.#logger.error("Checkpoint support disabled", { options });
+      return { success: false, reason: "DISABLED" };
+    }
+
+    const controller = new AbortController();
+    this.#abortControllers.set(runId, controller);
+
+    const $$ = $({ signal: controller.signal });
+
+    const shortCode = nanoid(8);
+    const imageRef = this.#getImageRef(projectRef, deploymentVersion, shortCode);
+    const exportLocation = this.#getExportLocation(projectRef, deploymentVersion, shortCode);
+
+    const cleanup = async () => {
+      if (this.#dockerMode) {
+        return;
+      }
+
+      try {
+        await $`rm ${exportLocation}`;
+        this.#logger.log("Deleted checkpoint archive", { exportLocation });
+
+        await $`buildah rmi ${imageRef}`;
+        this.#logger.log("Deleted checkpoint image", { imageRef });
+      } catch (error) {
+        this.#logger.error("Failure during checkpoint cleanup", { exportLocation, error });
+      }
+    };
+
+    try {
+      await this.chaosMonkey.call({ $: $$ });
+
+      this.#logger.log("Checkpointing:", { options });
+
+      const containterName = this.#getRunContainerName(runId);
+
+      // Create checkpoint (docker)
+      if (this.#dockerMode) {
+        try {
+          if (this.opts.forceSimulate || !this.#canCheckpoint) {
+            this.#logger.log("Simulating checkpoint");
+            this.#logger.debug(await $$`docker pause ${containterName}`);
+          } else {
+            if (this.simulateCheckpointFailure) {
+              if (performance.now() < this.simulateCheckpointFailureSeconds * 1000) {
+                this.#logger.error("Simulating checkpoint failure", { options });
+                throw new Error("SIMULATE_CHECKPOINT_FAILURE");
+              }
+            }
+
+            if (leaveRunning) {
+              this.#logger.debug(
+                await $$`docker checkpoint create --leave-running ${containterName} ${exportLocation}`
+              );
+            } else {
+              this.#logger.debug(
+                await $$`docker checkpoint create ${containterName} ${exportLocation}`
+              );
+            }
+          }
+        } catch (error) {
+          this.#logger.error("Failed while creating docker checkpoint", { exportLocation });
+          throw error;
+        }
+
+        this.#logger.log("checkpoint created:", {
+          runId,
+          location: exportLocation,
+        });
+
+        return {
+          success: true,
+          checkpoint: {
+            location: exportLocation,
+            docker: true,
+          },
+        };
+      }
+
+      // Create checkpoint (CRI)
+      if (!this.#canCheckpoint) {
+        this.#logger.error("No checkpoint support in kubernetes mode.");
+        return { success: false, reason: "SKIP_RETRYING" };
+      }
+
+      const containerId = this.#logger.debug(
+        // @ts-expect-error
+        await $$`crictl ps`
+          .pipeStdout($$({ stdin: "pipe" })`grep ${containterName}`)
+          .pipeStdout($$({ stdin: "pipe" })`cut -f1 ${"-d "}`)
+      );
+
+      if (!containerId.stdout) {
+        this.#logger.error("could not find container id", { options, containterName });
+        return { success: false, reason: "SKIP_RETRYING" };
+      }
+
+      const start = performance.now();
+
+      if (this.simulateCheckpointFailure) {
+        if (performance.now() < this.simulateCheckpointFailureSeconds * 1000) {
+          this.#logger.error("Simulating checkpoint failure", { options });
+          throw new Error("SIMULATE_CHECKPOINT_FAILURE");
+        }
+      }
+
+      // Create checkpoint
+      this.#logger.debug(await $$`crictl checkpoint --export=${exportLocation} ${containerId}`);
+      const postCheckpoint = performance.now();
+
+      // Print checkpoint size
+      const size = await getParsedFileSize(exportLocation);
+      this.#logger.log("checkpoint archive created", { size, options });
+
+      // Create image from checkpoint
+      const container = this.#logger.debug(await $$`buildah from scratch`);
+      const postFrom = performance.now();
+
+      this.#logger.debug(await $$`buildah add ${container} ${exportLocation} /`);
+      const postAdd = performance.now();
+
+      this.#logger.debug(
+        await $$`buildah config --annotation=io.kubernetes.cri-o.annotations.checkpoint.name=counter ${container}`
+      );
+      const postConfig = performance.now();
+
+      this.#logger.debug(await $$`buildah commit ${container} ${imageRef}`);
+      const postCommit = performance.now();
+
+      this.#logger.debug(await $$`buildah rm ${container}`);
+      const postRm = performance.now();
+
+      if (this.simulatePushFailure) {
+        if (performance.now() < this.simulatePushFailureSeconds * 1000) {
+          this.#logger.error("Simulating push failure", { options });
+          throw new Error("SIMULATE_PUSH_FAILURE");
+        }
+      }
+
+      // Push checkpoint image
+      this.#logger.debug(
+        await $$`buildah push --tls-verify=${String(this.registryTlsVerify)} ${imageRef}`
+      );
+      const postPush = performance.now();
+
+      const perf = {
+        "crictl checkpoint": postCheckpoint - start,
+        "buildah from": postFrom - postCheckpoint,
+        "buildah add": postAdd - postFrom,
+        "buildah config": postConfig - postAdd,
+        "buildah commit": postCommit - postConfig,
+        "buildah rm": postRm - postCommit,
+        "buildah push": postPush - postRm,
+      };
+
+      this.#logger.log("Checkpointed and pushed image to:", { location: imageRef, perf });
+
+      return {
+        success: true,
+        checkpoint: {
+          location: imageRef,
+          docker: false,
+        },
+      };
+    } catch (error) {
+      if (isExecaChildProcess(error)) {
+        if (error.isCanceled) {
+          this.#logger.error("Checkpoint canceled", { options, error });
+
+          return { success: false, reason: "CANCELED" };
+        }
+
+        this.#logger.error("Checkpoint command error", { options, error });
+
+        return { success: false, reason: "ERROR" };
+      }
+
+      this.#logger.error("Unhandled checkpoint error", { options, error });
+
+      return { success: false, reason: "ERROR" };
+    } finally {
+      this.#abortControllers.delete(runId);
+      await cleanup();
+    }
+  }
+
+  #failCheckpoint(runId: string, error: unknown) {
+    this.#failedCheckpoints.set(runId, error);
+  }
+
+  #clearFailedCheckpoint(runId: string) {
+    this.#failedCheckpoints.delete(runId);
+  }
+
+  #hasFailedCheckpoint(runId: string) {
+    return this.#failedCheckpoints.has(runId);
+  }
+
+  #getRunContainerName(suffix: string) {
+    return `task-run-${suffix}`;
+  }
+}

apps/docker-provider/package.json

@@ -18,8 +18,7 @@
   "dependencies": {
     "@trigger.dev/core": "workspace:*",
     "@trigger.dev/core-apps": "workspace:*",
-    "execa": "^8.0.1",
-    "socket.io-client": "^4.7.4"
+    "execa": "^8.0.1"
   },
   "devDependencies": {
     "@types/node": "^18.19.8",

apps/docker-provider/src/index.ts

@@ -1,14 +1,13 @@
 import { $, type ExecaChildProcess, execa } from "execa";
 import {
-  SimpleLogger,
-  TaskOperations,
   ProviderShell,
-  TaskOperationsRestoreOptions,
+  TaskOperations,
   TaskOperationsCreateOptions,
   TaskOperationsIndexOptions,
-  isExecaChildProcess,
-  testDockerCheckpoint,
-} from "@trigger.dev/core-apps";
+  TaskOperationsRestoreOptions,
+} from "@trigger.dev/core-apps/provider";
+import { SimpleLogger } from "@trigger.dev/core-apps/logger";
+import { isExecaChildProcess, testDockerCheckpoint } from "@trigger.dev/core-apps/checkpoints";
 import { setTimeout } from "node:timers/promises";
 import { PostStartCauses, PreStopCauses } from "@trigger.dev/core/v3";
 
@@ -54,13 +53,16 @@ class DockerTaskOperations implements TaskOperations {
   }
 
   #getInitReturn(canCheckpoint: boolean): TaskOperationsInitReturn {
-    this.#initialized = true;
     this.#canCheckpoint = canCheckpoint;
 
     if (canCheckpoint) {
-      logger.log("Full checkpoint support!");
+      if (!this.#initialized) {
+        logger.log("Full checkpoint support!");
+      }
     }
 
+    this.#initialized = true;
+
     const willSimulate = !canCheckpoint || this.opts.forceSimulate;
 
     if (willSimulate) {

apps/kubernetes-provider/package.json

@@ -19,8 +19,7 @@
     "@kubernetes/client-node": "^0.20.0",
     "@trigger.dev/core": "workspace:*",
     "@trigger.dev/core-apps": "workspace:*",
-    "p-queue": "^8.0.1",
-    "socket.io-client": "^4.7.4"
+    "p-queue": "^8.0.1"
   },
   "devDependencies": {
     "dotenv": "^16.4.2",

apps/kubernetes-provider/src/index.ts

@@ -1,12 +1,12 @@
 import * as k8s from "@kubernetes/client-node";
 import {
   ProviderShell,
-  SimpleLogger,
   TaskOperations,
   TaskOperationsCreateOptions,
   TaskOperationsIndexOptions,
   TaskOperationsRestoreOptions,
-} from "@trigger.dev/core-apps";
+} from "@trigger.dev/core-apps/provider";
+import { SimpleLogger } from "@trigger.dev/core-apps/logger";
 import {
   MachinePreset,
   PostStartCauses,

apps/kubernetes-provider/src/podCleaner.ts

@@ -1,5 +1,5 @@
 import * as k8s from "@kubernetes/client-node";
-import { SimpleLogger } from "@trigger.dev/core-apps";
+import { SimpleLogger } from "@trigger.dev/core-apps/logger";
 
 type PodCleanerOptions = {
   runtimeEnv: "local" | "kubernetes";

apps/kubernetes-provider/src/taskMonitor.ts

@@ -1,5 +1,5 @@
 import * as k8s from "@kubernetes/client-node";
-import { SimpleLogger } from "@trigger.dev/core-apps";
+import { SimpleLogger } from "@trigger.dev/core-apps/logger";
 import { EXIT_CODE_ALREADY_HANDLED, EXIT_CODE_CHILD_NONZERO } from "@trigger.dev/core-apps/process";
 import { setTimeout } from "timers/promises";
 import PQueue from "p-queue";
@@ -140,6 +140,9 @@ export class TaskMonitor {
     const exitCode = containerState.exitCode ?? -1;
 
     if (exitCode === EXIT_CODE_ALREADY_HANDLED) {
+      this.#logger.debug("Ignoring pod failure, already handled by worker", {
+        podName,
+      });
       return;
     }
 

apps/proxy/src/events/queueEvent.ts

@@ -24,6 +24,7 @@ export async function queueEvent(request: Request, env: Env): Promise<Response>
     const anyBody = await request.json();
     const body = SendEventBodySchema.safeParse(anyBody);
     if (!body.success) {
+      fromZodError(body.error);
       return json(
         { error: generateErrorMessage(body.error.issues) },
         {

apps/webapp/app/env.server.ts

@@ -151,7 +151,7 @@ const EnvironmentSchema = z.object({
   PROD_OTEL_LOG_EXPORT_TIMEOUT_MILLIS: z.string().default("30000"),
   PROD_OTEL_LOG_MAX_QUEUE_SIZE: z.string().default("512"),
 
-  RUNTIME_WAIT_THRESHOLD_IN_MS: z.coerce.number().int().default(30000),
+  CHECKPOINT_THRESHOLD_IN_MS: z.coerce.number().int().default(30000),
 
   // Internal OTEL environment variables
   INTERNAL_OTEL_TRACE_EXPORTER_URL: z.string().optional(),

apps/webapp/app/v3/environmentVariables/environmentVariablesRepository.server.ts

@@ -740,7 +740,7 @@ async function resolveBuiltInProdVariables(runtimeEnvironment: RuntimeEnvironmen
     },
     {
       key: "TRIGGER_RUNTIME_WAIT_THRESHOLD_IN_MS",
-      value: String(env.RUNTIME_WAIT_THRESHOLD_IN_MS),
+      value: String(env.CHECKPOINT_THRESHOLD_IN_MS),
     },
     {
       key: "TRIGGER_ORG_ID",

apps/webapp/app/v3/failedTaskRun.server.ts

@@ -9,14 +9,19 @@ import { BaseService } from "./services/baseService.server";
 const FAILABLE_TASK_RUN_STATUSES: TaskRunStatus[] = ["EXECUTING", "PENDING", "WAITING_FOR_DEPLOY"];
 
 export class FailedTaskRunService extends BaseService {
-  public async call(runFriendlyId: string, completion: TaskRunFailedExecutionResult) {
+  public async call(anyRunId: string, completion: TaskRunFailedExecutionResult) {
+    const isFriendlyId = anyRunId.startsWith("run_");
+
     const taskRun = await this._prisma.taskRun.findUnique({
-      where: { friendlyId: runFriendlyId },
+      where: {
+        friendlyId: isFriendlyId ? anyRunId : undefined,
+        id: !isFriendlyId ? anyRunId : undefined,
+      },
     });
 
     if (!taskRun) {
       logger.error("[FailedTaskRunService] Task run not found", {
-        runFriendlyId,
+        anyRunId,
         completion,
       });
 

apps/webapp/app/v3/handleSocketIo.server.ts

@@ -1,5 +1,6 @@
 import {
   ClientToSharedQueueMessages,
+  CoordinatorSocketData,
   CoordinatorToPlatformMessages,
   PlatformToCoordinatorMessages,
   PlatformToProviderMessages,
@@ -78,6 +79,7 @@ function createCoordinatorNamespace(io: Server) {
     authToken: env.COORDINATOR_SECRET,
     clientMessages: CoordinatorToPlatformMessages,
     serverMessages: PlatformToCoordinatorMessages,
+    socketData: CoordinatorSocketData,
     handlers: {
       READY_FOR_EXECUTION: async (message) => {
         const payload = await sharedQueueTasks.getLatestExecutionPayloadFromRun(
@@ -238,6 +240,45 @@ function createCoordinatorNamespace(io: Server) {
         }
       },
     },
+    onConnection: async (socket, handler, sender, logger) => {
+      if (socket.data.supportsDynamicConfig) {
+        socket.emit("DYNAMIC_CONFIG", {
+          version: "v1",
+          checkpointThresholdInMs: env.CHECKPOINT_THRESHOLD_IN_MS,
+        });
+      }
+    },
+    postAuth: async (socket, next, logger) => {
+      function setSocketDataFromHeader(
+        dataKey: keyof typeof socket.data,
+        headerName: string,
+        required: boolean = true
+      ) {
+        const value = socket.handshake.headers[headerName];
+
+        if (value) {
+          socket.data[dataKey] = Array.isArray(value) ? value[0] : value;
+          return;
+        }
+
+        if (required) {
+          logger.error("missing required header", { headerName });
+          throw new Error("missing header");
+        }
+      }
+
+      try {
+        setSocketDataFromHeader("supportsDynamicConfig", "x-supports-dynamic-config", false);
+      } catch (error) {
+        logger.error("setSocketDataFromHeader error", { error });
+        socket.disconnect(true);
+        return;
+      }
+
+      logger.debug("success", socket.data);
+
+      next();
+    },
   });
 
   return coordinator.namespace;

apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts

@@ -268,6 +268,7 @@ export class SharedQueueConsumer {
     // TODO: For every ACK, decide what should be done with the existing run and attempts. Make sure to check the current statuses first.
 
     switch (messageBody.data.type) {
+      // MARK: EXECUTE
       case "EXECUTE": {
         const existingTaskRun = await prisma.taskRun.findUnique({
           where: {
@@ -478,7 +479,7 @@ export class SharedQueueConsumer {
           ? lockedTaskRun.attempts[0].number + 1
           : 1;
 
-        const isRetry = nextAttemptNumber > 1;
+        const isRetry = lockedTaskRun.status === "WAITING_TO_RESUME" && nextAttemptNumber > 1;
 
         try {
           if (messageBody.data.checkpointEventId) {
@@ -493,6 +494,8 @@ export class SharedQueueConsumer {
               logger.error("Failed to restore checkpoint", {
                 queueMessage: message.data,
                 messageId: message.messageId,
+                runStatus: lockedTaskRun.status,
+                isRetry,
               });
 
               await this.#ackAndDoMoreWork(message.messageId);
@@ -503,8 +506,18 @@ export class SharedQueueConsumer {
           }
 
           if (!deployment.worker.supportsLazyAttempts) {
-            const service = new CreateTaskRunAttemptService();
-            await service.call(lockedTaskRun.friendlyId, undefined, false);
+            try {
+              const service = new CreateTaskRunAttemptService();
+              await service.call(lockedTaskRun.friendlyId, undefined, false);
+            } catch (error) {
+              logger.error("Failed to create task run attempt for outdate worker", {
+                error,
+                taskRun: lockedTaskRun.id,
+              });
+
+              await this.#ackAndDoMoreWork(message.messageId);
+              return;
+            }
           }
 
           if (isRetry) {
@@ -568,6 +581,7 @@ export class SharedQueueConsumer {
 
         break;
       }
+      // MARK: DEP RESUME
       // Resume after dependency completed with no remaining retries
       case "RESUME": {
         if (messageBody.data.checkpointEventId) {
@@ -728,6 +742,11 @@ export class SharedQueueConsumer {
         }
 
         try {
+          logger.debug("Broadcasting RESUME_AFTER_DEPENDENCY", {
+            runId: resumableAttempt.taskRunId,
+            attemptId: resumableAttempt.id,
+          });
+
           // The attempt should still be running so we can broadcast to all coordinators to resume immediately
           socketIo.coordinatorNamespace.emit("RESUME_AFTER_DEPENDENCY", {
             version: "v1",
@@ -752,6 +771,7 @@ export class SharedQueueConsumer {
 
         break;
       }
+      // MARK: DURATION RESUME
       // Resume after duration-based wait
       case "RESUME_AFTER_DURATION": {
         try {
@@ -785,6 +805,7 @@ export class SharedQueueConsumer {
 
         break;
       }
+      // MARK: FAIL
       // Fail for whatever reason, usually runs that have been resumed but stopped heartbeating
       case "FAIL": {
         const existingTaskRun = await prisma.taskRun.findUnique({
@@ -1122,6 +1143,11 @@ class SharedQueueTasks {
       },
       include: {
         lockedBy: true,
+        _count: {
+          select: {
+            attempts: true,
+          },
+        },
       },
     });
 
@@ -1143,6 +1169,7 @@ class SharedQueueTasks {
       runId: run.friendlyId,
       messageId: run.id,
       isTest: run.isTest,
+      attemptCount: run._count.attempts,
     } satisfies TaskRunExecutionLazyAttemptPayload;
   }
 

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

@@ -344,6 +344,7 @@ export class CompleteAttemptService extends BaseService {
     supportsLazyAttempts?: boolean
   ) {
     if (checkpointEventId || !supportsLazyAttempts) {
+      // Workers without lazy attempt support always need to go through the queue, which is where the attempt is created
       // We have to replace a potential RESUME with EXECUTE to correctly retry the attempt
       return await marqs?.replaceMessage(
         run.id,
@@ -355,8 +356,9 @@ export class CompleteAttemptService extends BaseService {
         retryTimestamp
       );
     } else {
-      // There's no checkpoint so the worker is still running and waiting for a retry message
-      // It supports lazy attempts so we can bypass the queue and send the message directly to the worker
+      // There's no checkpoint and the worker supports lazy attempts
+      // This means the worker is still running and waiting for a retry message
+      // It supports lazy attempts so we can bypass the queue and send the message directly to it
       RetryAttemptService.enqueue(run.id, this._prisma, new Date(retryTimestamp));
     }
   }

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

@@ -7,6 +7,8 @@ import { BaseService, ServiceValidationError } from "./baseService.server";
 import { TaskRun, TaskRunAttempt } from "@trigger.dev/database";
 import { machinePresetFromConfig } from "../machinePresets.server";
 import { workerQueue } from "~/services/worker.server";
+import { MAX_TASK_RUN_ATTEMPTS } from "~/consts";
+import { CrashTaskRunService } from "./crashTaskRun.server";
 
 export class CreateTaskRunAttemptService extends BaseService {
   public async call(
@@ -93,6 +95,17 @@ export class CreateTaskRunAttemptService extends BaseService {
 
       const nextAttemptNumber = taskRun.attempts[0] ? taskRun.attempts[0].number + 1 : 1;
 
+      if (nextAttemptNumber > MAX_TASK_RUN_ATTEMPTS) {
+        const service = new CrashTaskRunService(this._prisma);
+        await service.call(taskRun.id, {
+          reason: taskRun.lockedBy.worker.supportsLazyAttempts
+            ? "Max attempts reached."
+            : "Max attempts reached. Please upgrade your CLI and SDK.",
+        });
+
+        throw new ServiceValidationError("Max attempts reached", 400);
+      }
+
       const taskRunAttempt = await $transaction(this._prisma, async (tx) => {
         const taskRunAttempt = await tx.taskRunAttempt.create({
           data: {

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

@@ -12,6 +12,7 @@ import { socketIo } from "../handleSocketIo.server";
 import { SharedQueueMessageBody, sharedQueueTasks } from "../marqs/sharedQueueConsumer.server";
 import { BaseService } from "./baseService.server";
 import { TaskRunAttempt } from "@trigger.dev/database";
+import { isFinalRunStatus } from "../taskStatus";
 
 export class ResumeAttemptService extends BaseService {
   public async call(
@@ -80,10 +81,11 @@ export class ResumeAttemptService extends BaseService {
         return;
       }
 
-      if (attempt.taskRun.status !== "WAITING_TO_RESUME") {
+      if (isFinalRunStatus(attempt.taskRun.status)) {
         logger.error("Run is not resumable", {
           attemptId: attempt.id,
           runId: attempt.taskRunId,
+          status: attempt.taskRun.status,
         });
         return;
       }

packages/cli-v3/package.json

@@ -117,7 +117,6 @@
     "react-error-boundary": "^4.0.12",
     "semver": "^7.5.0",
     "simple-git": "^3.19.0",
-    "socket.io-client": "^4.7.4",
     "source-map-support": "^0.5.21",
     "terminal-link": "^3.0.0",
     "tiny-invariant": "^1.2.0",

packages/cli-v3/src/commands/deploy.ts

@@ -1412,10 +1412,10 @@ async function resolveEnvironmentVariables(
   const projectConfig = config.config;
 
   return await tracer.startActiveSpan("resolveEnvironmentVariables", async (span) => {
-    try {
-      const $spinner = spinner();
-      $spinner.start("Resolving environment variables");
+    const $spinner = spinner();
+    $spinner.start("Resolving environment variables");
 
+    try {
       let processEnv: Record<string, string | undefined> = {
         ...process.env,
       };
@@ -1473,10 +1473,14 @@ async function resolveEnvironmentVariables(
           $spinner.stop("No environment variables to sync");
           return;
         }
+      } else {
+        $spinner.stop("No environment variables to sync");
       }
 
       $spinner.stop("Environment variables resolved");
     } catch (e) {
+      $spinner.stop("Failed to resolve environment variables");
+
       recordSpanException(span, e);
 
       throw e;

packages/cli-v3/src/commands/dev.tsx

@@ -298,6 +298,7 @@ function useDev({
     websocket.addEventListener("error", (event) => {});
 
     // This is the deprecated task heart beat that uses the friendly attempt ID
+    // It will only be used if the worker does not support lazy attempts
     backgroundWorkerCoordinator.onWorkerTaskHeartbeat.attach(
       async ({ worker, backgroundWorkerId, id }) => {
         await sender.send("BACKGROUND_WORKER_MESSAGE", {

packages/cli-v3/src/workers/prod/backgroundWorker.ts

@@ -14,7 +14,6 @@ import {
   TaskRunExecutionLazyAttemptPayload,
   TaskRunExecutionPayload,
   TaskRunExecutionResult,
-  WaitReason,
   correctErrorStackTrace,
 } from "@trigger.dev/core/v3";
 import { ZodIpcConnection } from "@trigger.dev/core/v3/zodIpc";
@@ -39,6 +38,19 @@ type BackgroundWorkerParams = {
   debugOtel?: boolean;
 };
 
+export type OnWaitForDurationMessage = InferSocketMessageSchema<
+  typeof ProdChildToWorkerMessages,
+  "WAIT_FOR_DURATION"
+>;
+export type OnWaitForTaskMessage = InferSocketMessageSchema<
+  typeof ProdChildToWorkerMessages,
+  "WAIT_FOR_TASK"
+>;
+export type OnWaitForBatchMessage = InferSocketMessageSchema<
+  typeof ProdChildToWorkerMessages,
+  "WAIT_FOR_BATCH"
+>;
+
 export class ProdBackgroundWorker {
   private _initialized: boolean = false;
 
@@ -48,21 +60,9 @@ export class ProdBackgroundWorker {
   public onTaskHeartbeat: Evt<string> = new Evt();
   public onTaskRunHeartbeat: Evt<string> = new Evt();
 
-  public onWaitForBatch: Evt<
-    InferSocketMessageSchema<typeof ProdChildToWorkerMessages, "WAIT_FOR_BATCH">
-  > = new Evt();
-  public onWaitForDuration: Evt<
-    InferSocketMessageSchema<typeof ProdChildToWorkerMessages, "WAIT_FOR_DURATION">
-  > = new Evt();
-  public onWaitForTask: Evt<
-    InferSocketMessageSchema<typeof ProdChildToWorkerMessages, "WAIT_FOR_TASK">
-  > = new Evt();
-
-  public preCheckpointNotification = Evt.create<{ willCheckpointAndRestore: boolean }>();
-  public checkpointCanceledNotification = Evt.create<{ checkpointCanceled: boolean }>();
-
-  public onReadyForCheckpoint = Evt.create<{ version?: "v1" }>();
-  public onCancelCheckpoint = Evt.create<{ version?: "v1" | "v2"; reason?: WaitReason }>();
+  public onWaitForDuration: Evt<OnWaitForDurationMessage> = new Evt();
+  public onWaitForTask: Evt<OnWaitForTaskMessage> = new Evt();
+  public onWaitForBatch: Evt<OnWaitForBatchMessage> = new Evt();
 
   public onCreateTaskRunAttempt = Evt.create<{ version?: "v1"; runId: string }>();
   public attemptCreatedNotification = Evt.create<
@@ -133,7 +133,12 @@ export class ProdBackgroundWorker {
   }
 
   async flushTelemetry() {
+    console.log("Flushing telemetry");
+    const start = performance.now();
+
     await this._taskRunProcess?.cleanup(false);
+
+    console.log("Flushed telemetry", { duration: performance.now() - start });
   }
 
   async initialize(options?: { env?: Record<string, string> }) {
@@ -302,22 +307,6 @@ export class ProdBackgroundWorker {
       this.onWaitForTask.post(message);
     });
 
-    taskRunProcess.onReadyForCheckpoint.attach((message) => {
-      this.onReadyForCheckpoint.post(message);
-    });
-
-    taskRunProcess.onCancelCheckpoint.attach((message) => {
-      this.onCancelCheckpoint.post(message);
-    });
-
-    // Notify down the chain
-    this.preCheckpointNotification.attach((message) => {
-      taskRunProcess.preCheckpointNotification.post(message);
-    });
-    this.checkpointCanceledNotification.attach((message) => {
-      taskRunProcess.checkpointCanceledNotification.post(message);
-    });
-
     await taskRunProcess.initialize();
 
     this._taskRunProcess = taskRunProcess;
@@ -373,6 +362,8 @@ export class ProdBackgroundWorker {
     kill = false,
     initialSignal: number | NodeJS.Signals = "SIGTERM"
   ) {
+    console.log("Trying graceful exit", { kill, initialSignal });
+
     try {
       const initialExit = taskRunProcess.onExit.waitFor(5_000);
 
@@ -389,6 +380,8 @@ export class ProdBackgroundWorker {
   }
 
   async #tryForcefulExit(taskRunProcess: TaskRunProcess) {
+    console.log("Trying forceful exit");
+
     try {
       const forcedKill = taskRunProcess.onExit.waitFor(5_000);
       taskRunProcess.kill("SIGKILL");
@@ -524,19 +517,24 @@ export class ProdBackgroundWorker {
     let execution: ProdTaskRunExecution;
 
     try {
+      const start = performance.now();
+
       // ..and wait for response
-      const attemptCreated = await this.attemptCreatedNotification.waitFor(30_000);
+      const attemptCreated = await this.attemptCreatedNotification.waitFor(120_000);
 
       if (!attemptCreated.success) {
-        throw new Error(
-          `Failed to create attempt${attemptCreated.reason ? `: ${attemptCreated.reason}` : ""}`
-        );
+        throw new Error(`${attemptCreated.reason ?? "Unknown error"}`);
       }
 
+      console.log("Attempt created", {
+        number: attemptCreated.execution.attempt.number,
+        duration: performance.now() - start,
+      });
+
       execution = attemptCreated.execution;
     } catch (error) {
       console.error("Error while creating attempt", error);
-      throw new Error(`Failed to create task run attempt: ${error}`);
+      throw new Error(`Failed to create attempt: ${error}`);
     }
 
     const completion = await this.executeTaskRun(
@@ -590,21 +588,11 @@ class TaskRunProcess {
     new Evt();
   public onIsBeingKilled: Evt<TaskRunProcess> = new Evt();
 
-  public onWaitForBatch: Evt<
-    InferSocketMessageSchema<typeof ProdChildToWorkerMessages, "WAIT_FOR_BATCH">
-  > = new Evt();
-  public onWaitForDuration: Evt<
-    InferSocketMessageSchema<typeof ProdChildToWorkerMessages, "WAIT_FOR_DURATION">
-  > = new Evt();
-  public onWaitForTask: Evt<
-    InferSocketMessageSchema<typeof ProdChildToWorkerMessages, "WAIT_FOR_TASK">
-  > = new Evt();
+  public onWaitForDuration: Evt<OnWaitForDurationMessage> = new Evt();
+  public onWaitForTask: Evt<OnWaitForTaskMessage> = new Evt();
+  public onWaitForBatch: Evt<OnWaitForBatchMessage> = new Evt();
 
   public preCheckpointNotification = Evt.create<{ willCheckpointAndRestore: boolean }>();
-  public checkpointCanceledNotification = Evt.create<{ checkpointCanceled: boolean }>();
-
-  public onReadyForCheckpoint = Evt.create<{ version?: "v1" }>();
-  public onCancelCheckpoint = Evt.create<{ version?: "v1" | "v2"; reason?: WaitReason }>();
 
   constructor(
     private runId: string,
@@ -664,6 +652,10 @@ class TaskRunProcess {
           if (this.messageId) {
             this.onTaskRunHeartbeat.post(this.messageId);
           } else {
+            console.error(
+              "No message id for task heartbeat, falling back to (deprecated) attempt heartbeat",
+              { id: message.id }
+            );
             this.onTaskHeartbeat.post(message.id);
           }
         },
@@ -675,55 +667,7 @@ class TaskRunProcess {
           this.onWaitForBatch.post(message);
         },
         WAIT_FOR_DURATION: async (message) => {
-          // Post to coordinator
           this.onWaitForDuration.post(message);
-
-          try {
-            // ..and wait for response
-            const { willCheckpointAndRestore } = await this.preCheckpointNotification.waitFor(
-              30_000
-            );
-
-            return {
-              willCheckpointAndRestore,
-            };
-          } catch (error) {
-            console.error("Error while waiting for pre-checkpoint notification", error);
-
-            // Assume we won't get checkpointed
-            return {
-              willCheckpointAndRestore: false,
-            };
-          }
-        },
-        READY_FOR_CHECKPOINT: async (message) => {
-          this.onReadyForCheckpoint.post(message);
-        },
-        CANCEL_CHECKPOINT: async (message) => {
-          const version = "v2";
-
-          // Post to coordinator
-          this.onCancelCheckpoint.post(message);
-
-          try {
-            // ..and wait for response
-            const { checkpointCanceled } = await this.checkpointCanceledNotification.waitFor(
-              30_000
-            );
-
-            return {
-              version,
-              checkpointCanceled,
-            };
-          } catch (error) {
-            console.error("Error while waiting for checkpoint cancellation", error);
-
-            // Assume it's been canceled
-            return {
-              version,
-              checkpointCanceled: true,
-            };
-          }
         },
       },
     });
@@ -764,14 +708,21 @@ class TaskRunProcess {
       realChildPid: this._child?.pid,
     });
 
-    await this._ipc?.sendWithAck(
-      "CLEANUP",
-      {
-        flush: true,
-        kill: killParentProcess,
-      },
-      30_000
-    );
+    try {
+      await this._ipc?.sendWithAck(
+        "CLEANUP",
+        {
+          flush: true,
+          kill: killParentProcess,
+        },
+        30_000
+      );
+    } catch (error) {
+      console.error("Error while cleaning up task run process", error);
+      if (killParentProcess) {
+        process.exit(0);
+      }
+    }
 
     if (killChildProcess) {
       this._gracefulExitTimeoutElapsed = true;
@@ -815,21 +766,34 @@ class TaskRunProcess {
 
   taskRunCompletedNotification(completion: TaskRunExecutionResult) {
     if (!completion.ok && typeof completion.retry !== "undefined") {
+      console.error(
+        "Task run completed with error and wants to retry, won't send task run completed notification"
+      );
       return;
     }
 
-    if (this._child?.connected && !this._isBeingKilled && !this._child.killed) {
-      this._ipc?.send("TASK_RUN_COMPLETED_NOTIFICATION", {
-        version: "v2",
-        completion,
-      });
+    if (!this._child?.connected || this._isBeingKilled || this._child.killed) {
+      console.error(
+        "Child process not connected or being killed, can't send task run completed notification"
+      );
+      return;
     }
+
+    this._ipc?.send("TASK_RUN_COMPLETED_NOTIFICATION", {
+      version: "v2",
+      completion,
+    });
   }
 
   waitCompletedNotification() {
-    if (this._child?.connected && !this._isBeingKilled && !this._child.killed) {
-      this._ipc?.send("WAIT_COMPLETED_NOTIFICATION", {});
+    if (!this._child?.connected || this._isBeingKilled || this._child.killed) {
+      console.error(
+        "Child process not connected or being killed, can't send wait completed notification"
+      );
+      return;
     }
+
+    this._ipc?.send("WAIT_COMPLETED_NOTIFICATION", {});
   }
 
   async #handleExit(code: number | null, signal: NodeJS.Signals | null) {

packages/cli-v3/src/workers/prod/worker-facade.ts

@@ -292,6 +292,7 @@ async function asyncHeartbeat(initialDelayInSeconds: number = 30, intervalInSeco
     while (true) {
       if (_isRunning && _execution) {
         try {
+          // The attempt ID will only be used to heartbeat if the message (run) ID isn't set on the TaskRunProcess
           await zodIpc.send("TASK_HEARTBEAT", { id: _execution.attempt.id });
         } catch (err) {
           console.error("Failed to send HEARTBEAT message", err);

packages/core-apps/package.json

@@ -28,7 +28,6 @@
     "@trigger.dev/core": "workspace:*",
     "@trigger.dev/tsconfig": "workspace:*",
     "@types/node": "18",
-    "socket.io-client": "^4.7.4",
     "typescript": "^5.3.0"
   },
   "engines": {

packages/core-apps/src/backoff.ts

@@ -1,3 +1,5 @@
+import { setTimeout as timeout } from "node:timers/promises";
+
 type ExponentialBackoffType = "NoJitter" | "FullJitter" | "EqualJitter";
 
 type ExponentialBackoffOptions = {
@@ -16,6 +18,26 @@ class StopRetrying extends Error {
   }
 }
 
+class AttemptTimeout extends Error {
+  constructor(message?: string) {
+    super(message);
+    this.name = "AttemptTimeout";
+  }
+}
+
+class RetryLimitExceeded extends Error {
+  constructor(message?: string) {
+    super(message);
+    this.name = "RetryLimitExceeded";
+  }
+}
+
+type YieldType<T> = T extends AsyncGenerator<infer Y, any, any> ? Y : never;
+
+/**
+ * Exponential backoff helper class
+ * - All time units in seconds unless otherwise specified
+ */
 export class ExponentialBackoff {
   #retries: number = 0;
 
@@ -41,64 +63,44 @@ export class ExponentialBackoff {
     this.#maxElapsed = opts.maxElapsed ?? Infinity;
   }
 
-  #clone() {
-    return new ExponentialBackoff(this.#type, {
-      base: this.#base,
-      factor: this.#factor,
-      min: this.#min,
-      max: this.#max,
-      maxRetries: this.#maxRetries,
-      maxElapsed: this.#maxElapsed,
+  #clone(type?: ExponentialBackoffType, opts: Partial<ExponentialBackoffOptions> = {}) {
+    return new ExponentialBackoff(type ?? this.#type, {
+      base: opts.base ?? this.#base,
+      factor: opts.factor ?? this.#factor,
+      min: opts.min ?? this.#min,
+      max: opts.max ?? this.#max,
+      maxRetries: opts.maxRetries ?? this.#maxRetries,
+      maxElapsed: opts.maxElapsed ?? this.#maxElapsed,
     });
   }
 
   type(type?: ExponentialBackoffType) {
-    if (typeof type !== "undefined") {
-      this.#type = type;
-    }
-    return this.#clone();
+    return this.#clone(type);
   }
 
   base(base?: number) {
-    if (typeof base !== "undefined") {
-      this.#base = base;
-    }
-    return this.#clone();
+    return this.#clone(undefined, { base });
   }
 
   factor(factor?: number) {
-    if (typeof factor !== "undefined") {
-      this.#factor = factor;
-    }
-    return this.#clone();
+    return this.#clone(undefined, { factor });
   }
 
   min(min?: number) {
-    if (typeof min !== "undefined") {
-      this.#min = min;
-    }
-    return this.#clone();
+    return this.#clone(undefined, { min });
   }
 
   max(max?: number) {
-    if (typeof max !== "undefined") {
-      this.#max = max;
-    }
-    return this.#clone();
+    return this.#clone(undefined, { max });
   }
 
   maxRetries(maxRetries?: number) {
-    if (typeof maxRetries !== "undefined") {
-      this.#maxRetries = maxRetries;
-    }
-    return this.#clone();
+    return this.#clone(undefined, { maxRetries });
   }
 
+  // TODO: With .execute(), should this also include the time it takes to execute the callback?
   maxElapsed(maxElapsed?: number) {
-    if (typeof maxElapsed !== "undefined") {
-      this.#maxElapsed = maxElapsed;
-    }
-    return this.#clone();
+    return this.#clone(undefined, { maxElapsed });
   }
 
   retries(retries?: number) {
@@ -145,6 +147,7 @@ export class ExponentialBackoff {
     yield* this.retryAsync();
   }
 
+  /** Returns the delay for the current retry in seconds. */
   delay(retries: number = this.#retries, jitter: boolean = true) {
     if (retries > this.#maxRetries) {
       console.error(
@@ -184,13 +187,31 @@ export class ExponentialBackoff {
       }
     }
 
-    delay = Math.min(delay, this.#max);
-    delay = Math.max(delay, this.#min);
+    // If min/max override the delay, jitter with 20% while respecting min/max
+    if (delay < this.#min) {
+      delay = this.#min + Math.random() * (this.#min * 0.2);
+    }
+    if (delay > this.#max) {
+      delay = this.#max - Math.random() * (this.#max * 0.2);
+    }
+
     delay = Math.round(delay);
 
     return delay;
   }
 
+  /** Waits with the appropriate delay for the current retry. */
+  async wait(retries: number = this.#retries, jitter: boolean = true) {
+    if (retries > this.#maxRetries) {
+      console.error(`Retry limit exceeded: ${retries} > ${this.#maxRetries}`);
+      throw new RetryLimitExceeded();
+    }
+
+    const delay = this.delay(retries, jitter);
+
+    return await timeout(delay * 1000);
+  }
+
   elapsed(retries: number = this.#retries, jitter: boolean = true) {
     let elapsed = 0;
 
@@ -243,5 +264,98 @@ export class ExponentialBackoff {
     throw new StopRetrying();
   }
 
+  get state() {
+    return {
+      retries: this.#retries,
+      type: this.#type,
+      base: this.#base,
+      factor: this.#factor,
+      min: this.#min,
+      max: this.#max,
+      maxRetries: this.#maxRetries,
+      maxElapsed: this.#maxElapsed,
+    };
+  }
+
+  async execute<T>(
+    callback: (
+      iteratorReturn: YieldType<ReturnType<ExponentialBackoff["retryAsync"]>> & {
+        elapsedMs: number;
+      }
+    ) => Promise<T>,
+    { attemptTimeoutMs = 0 }: { attemptTimeoutMs?: number } = {}
+  ): Promise<
+    | { success: true; result: T }
+    | { success: false; error?: unknown; cause: "StopRetrying" | "Timeout" | "MaxRetries" }
+  > {
+    let elapsedMs = 0;
+    let finalError: unknown = undefined;
+
+    for await (const { delay, retry } of this) {
+      const start = Date.now();
+
+      if (retry > 0) {
+        console.log(`Retrying in ${delay.milliseconds}ms`);
+        await timeout(delay.milliseconds);
+      }
+
+      let attemptTimeout: NodeJS.Timeout | undefined = undefined;
+
+      try {
+        const result = await new Promise<T>(async (resolve, reject) => {
+          if (attemptTimeoutMs > 0) {
+            attemptTimeout = setTimeout(() => {
+              reject(new AttemptTimeout());
+            }, attemptTimeoutMs);
+          }
+
+          try {
+            const callbackResult = await callback({ delay, retry, elapsedMs });
+
+            resolve(callbackResult);
+          } catch (error) {
+            reject(error);
+          }
+        });
+
+        return {
+          success: true,
+          result,
+        };
+      } catch (error) {
+        finalError = error;
+
+        if (error instanceof StopRetrying) {
+          return {
+            success: false,
+            cause: "StopRetrying",
+            error: error.message,
+          };
+        }
+
+        if (error instanceof AttemptTimeout) {
+          continue;
+        }
+      } finally {
+        elapsedMs += Date.now() - start;
+        clearTimeout(attemptTimeout);
+      }
+    }
+
+    if (finalError instanceof AttemptTimeout) {
+      return {
+        success: false,
+        cause: "Timeout",
+      };
+    } else {
+      return {
+        success: false,
+        cause: "MaxRetries",
+        error: finalError,
+      };
+    }
+  }
+
+  static RetryLimitExceeded = RetryLimitExceeded;
   static StopRetrying = StopRetrying;
 }

packages/core-apps/src/index.ts

@@ -1,4 +0,0 @@
-export * from "./http";
-export * from "./logger";
-export * from "./provider";
-export * from "./checkpoints";

packages/core/package.json

@@ -159,7 +159,7 @@
     "@opentelemetry/sdk-trace-node": "^1.22.0",
     "@opentelemetry/semantic-conventions": "^1.22.0",
     "humanize-duration": "^3.27.3",
-    "socket.io-client": "4.7.4",
+    "socket.io-client": "4.7.5",
     "superjson": "^2.2.1",
     "ulidx": "^2.2.1",
     "zod": "3.22.3",

packages/core/src/v3/runtime/prodRuntimeManager.ts

@@ -1,14 +1,11 @@
 import { clock } from "../clock-api";
-import { logger } from "../logger-api";
 import {
   BatchTaskRunExecutionResult,
   ProdChildToWorkerMessages,
   ProdWorkerToChildMessages,
   TaskRunContext,
-  TaskRunExecution,
   TaskRunExecutionResult,
 } from "../schemas";
-import { checkpointSafeTimeout, unboundedTimeout } from "../utils/timers";
 import { ZodIpcConnection } from "../zodIpc";
 import { RuntimeManager } from "./manager";
 
@@ -24,9 +21,7 @@ export class ProdRuntimeManager implements RuntimeManager {
     { resolve: (value: BatchTaskRunExecutionResult) => void; reject: (err?: any) => void }
   > = new Map();
 
-  _waitForDuration:
-    | { resolve: (value: "external") => void; reject: (err?: any) => void }
-    | undefined;
+  _waitForDuration: { resolve: (value: void) => void; reject: (err?: any) => void } | undefined;
 
   constructor(
     private ipc: ZodIpcConnection<
@@ -43,64 +38,17 @@ export class ProdRuntimeManager implements RuntimeManager {
   async waitForDuration(ms: number): Promise<void> {
     const now = Date.now();
 
-    const internalTimeout = unboundedTimeout(ms, "internal" as const);
-    const checkpointSafeInternalTimeout = checkpointSafeTimeout(ms);
-
-    if (ms <= this.waitThresholdInMs) {
-      await internalTimeout;
-      return;
-    }
-
-    const externalResume = new Promise<"external">((resolve, reject) => {
+    const resume = new Promise<void>((resolve, reject) => {
       this._waitForDuration = { resolve, reject };
     });
 
-    const { willCheckpointAndRestore } = await this.ipc.sendWithAck(
-      "WAIT_FOR_DURATION",
-      {
-        ms,
-        now,
-      },
-      31_000
-    );
+    await this.ipc.send("WAIT_FOR_DURATION", {
+      ms,
+      now,
+      waitThresholdInMs: this.waitThresholdInMs,
+    });
 
-    if (!willCheckpointAndRestore) {
-      await internalTimeout;
-      return;
-    }
-
-    this.ipc.send("READY_FOR_CHECKPOINT", {});
-
-    // internalTimeout acts as a backup and will be accurate if the checkpoint never happens
-    // checkpointSafeInternalTimeout is accurate even after non-simulated restores
-    await Promise.race([internalTimeout, checkpointSafeInternalTimeout]);
-
-    // Resets the clock to the current time
-    clock.reset();
-
-    try {
-      // The coordinator should cancel any in-progress checkpoints
-      const { checkpointCanceled, version } = await this.ipc.sendWithAck(
-        "CANCEL_CHECKPOINT",
-        {
-          version: "v2",
-          reason: "WAIT_FOR_DURATION",
-        },
-        31_000
-      );
-
-      if (checkpointCanceled) {
-        // There won't be a checkpoint or external resume and we've already completed our internal timeout
-        return;
-      }
-    } catch (error) {
-      // If the cancellation times out, we will proceed as if the checkpoint was canceled
-      logger.debug("Checkpoint cancellation timed out", { error });
-      return;
-    }
-
-    // No checkpoint was canceled, so we were checkpointed. We need to wait for the external resume message.
-    await externalResume;
+    await resume;
   }
 
   resumeAfterDuration(): void {
@@ -111,7 +59,7 @@ export class ProdRuntimeManager implements RuntimeManager {
     // Resets the clock to the current time
     clock.reset();
 
-    this._waitForDuration.resolve("external");
+    this._waitForDuration.resolve();
     this._waitForDuration = undefined;
   }
 

packages/core/src/v3/schemas/messages.ts

@@ -212,47 +212,15 @@ export const ProdChildToWorkerMessages = {
       id: z.string(),
     }),
   },
-  TASK_RUN_HEARTBEAT: {
-    message: z.object({
-      version: z.literal("v1").default("v1"),
-      id: z.string(),
-    }),
-  },
   READY_TO_DISPOSE: {
     message: z.undefined(),
   },
-  READY_FOR_CHECKPOINT: {
-    message: z.object({
-      version: z.literal("v1").default("v1"),
-    }),
-  },
-  CANCEL_CHECKPOINT: {
-    message: z
-      .discriminatedUnion("version", [
-        z.object({
-          version: z.literal("v1"),
-        }),
-        z.object({
-          version: z.literal("v2"),
-          reason: WaitReason.optional(),
-        }),
-      ])
-      .default({ version: "v1" }),
-    callback: z.object({
-      // TODO: Figure out how best to handle callback schema parsing in zod IPC
-      version: z.literal("v2") /* .default("v2") */,
-      checkpointCanceled: z.boolean(),
-      reason: WaitReason.optional(),
-    }),
-  },
   WAIT_FOR_DURATION: {
     message: z.object({
       version: z.literal("v1").default("v1"),
       ms: z.number(),
       now: z.number(),
-    }),
-    callback: z.object({
-      willCheckpointAndRestore: z.boolean(),
+      waitThresholdInMs: z.number(),
     }),
   },
   WAIT_FOR_TASK: {
@@ -457,6 +425,7 @@ export const CoordinatorToPlatformMessages = {
       }),
     ]),
   },
+  // Deprecated: Only workers without lazy attempt support will use this
   READY_FOR_EXECUTION: {
     message: z.object({
       version: z.literal("v1").default("v1"),
@@ -624,6 +593,12 @@ export const PlatformToCoordinatorMessages = {
       runId: z.string(),
     }),
   },
+  DYNAMIC_CONFIG: {
+    message: z.object({
+      version: z.literal("v1").default("v1"),
+      checkpointThresholdInMs: z.number(),
+    }),
+  },
 };
 
 export const ClientToSharedQueueMessages = {
@@ -672,10 +647,9 @@ const IndexTasksMessage = z.object({
 });
 
 export const ProdWorkerToCoordinatorMessages = {
-  LOG: {
+  TEST: {
     message: z.object({
       version: z.literal("v1").default("v1"),
-      text: z.string(),
     }),
     callback: z.void(),
   },
@@ -698,6 +672,7 @@ export const ProdWorkerToCoordinatorMessages = {
       }),
     ]),
   },
+  // Deprecated: Only workers without lazy attempt support will use this
   READY_FOR_EXECUTION: {
     message: z.object({
       version: z.literal("v1").default("v1"),
@@ -784,7 +759,7 @@ export const ProdWorkerToCoordinatorMessages = {
   },
   WAIT_FOR_TASK: {
     message: z.object({
-      version: z.literal("v1").default("v1"),
+      version: z.enum(["v1", "v2"]).default("v1"),
       friendlyId: z.string(),
       // This is the attempt that is waiting
       attemptFriendlyId: z.string(),
@@ -795,7 +770,7 @@ export const ProdWorkerToCoordinatorMessages = {
   },
   WAIT_FOR_BATCH: {
     message: z.object({
-      version: z.literal("v1").default("v1"),
+      version: z.enum(["v1", "v2"]).default("v1"),
       batchFriendlyId: z.string(),
       runFriendlyIds: z.string().array(),
       // This is the attempt that is waiting
@@ -843,6 +818,12 @@ export const ProdWorkerToCoordinatorMessages = {
       }),
     }),
   },
+  SET_STATE: {
+    message: z.object({
+      version: z.literal("v1").default("v1"),
+      attemptFriendlyId: z.string().optional(),
+    }),
+  },
 };
 
 // TODO: The coordinator can only safely use v1 worker messages, higher versions will need a new flag, e.g. SUPPORTS_VERSIONED_MESSAGES
@@ -861,6 +842,7 @@ export const CoordinatorToProdWorkerMessages = {
       attemptId: z.string(),
     }),
   },
+  // Deprecated: Only workers without lazy attempt support will use this
   EXECUTE_TASK_RUN: {
     message: z.object({
       version: z.literal("v1").default("v1"),
@@ -908,3 +890,7 @@ export const ProdWorkerSocketData = z.object({
   deploymentId: z.string(),
   deploymentVersion: z.string(),
 });
+
+export const CoordinatorSocketData = z.object({
+  supportsDynamicConfig: z.string().optional(),
+});

packages/core/src/v3/schemas/schemas.ts

@@ -220,6 +220,7 @@ export type WaitReason = z.infer<typeof WaitReason>;
 
 export const TaskRunExecutionLazyAttemptPayload = z.object({
   runId: z.string(),
+  attemptCount: z.number().optional(),
   messageId: z.string(),
   isTest: z.boolean(),
   traceContext: z.record(z.unknown()),

packages/core/src/v3/zodSocket.ts

@@ -1,4 +1,4 @@
-import type { Socket } from "socket.io-client";
+import type { ManagerOptions, Socket, SocketOptions } from "socket.io-client";
 import { io } from "socket.io-client";
 import { ZodError, z } from "zod";
 import { EventEmitterLike, ZodMessageValueSchema } from "./zodMessageHandler";
@@ -67,6 +67,7 @@ export type ZodSocketMessageHandlerOptions<TMessageCatalog extends ZodSocketMess
   {
     schema: TMessageCatalog;
     handlers?: ZodSocketMessageHandlers<TMessageCatalog>;
+    logger?: StructuredLogger;
   };
 
 type MessageFromSocketSchema<
@@ -90,42 +91,69 @@ const messageSchema = z.object({
 export class ZodSocketMessageHandler<TRPCCatalog extends ZodSocketMessageCatalogSchema> {
   #schema: TRPCCatalog;
   #handlers: ZodSocketMessageHandlers<TRPCCatalog> | undefined;
+  #logger: StructuredLogger;
 
   constructor(options: ZodSocketMessageHandlerOptions<TRPCCatalog>) {
     this.#schema = options.schema;
     this.#handlers = options.handlers;
+    this.#logger =
+      options.logger ?? new SimpleStructuredLogger("socket-message-handler", LogLevel.info);
   }
 
   public async handleMessage(message: unknown) {
-    const parsedMessage = this.parseMessage(message);
+    const parseResult = this.parseMessage(message);
+
+    if (!parseResult.success) {
+      this.#logger.error("Failed to parse message, skipping handler", {
+        rawMessage: message,
+        error: parseResult.reason,
+      });
+      return;
+    }
 
     if (!this.#handlers) {
       throw new Error("No handlers provided");
     }
 
-    const handler = this.#handlers[parsedMessage.type];
+    const { type, payload } = parseResult.data;
+
+    const handler = this.#handlers[type];
 
     if (!handler) {
-      console.error(`No handler for message type: ${String(parsedMessage.type)}`);
+      console.error(`No handler for message type: ${String(type)}`);
       return;
     }
 
-    const ack = await handler(parsedMessage.payload);
+    const ack = await handler(payload);
 
     return ack;
   }
 
-  public parseMessage(message: unknown): MessagesFromSocketCatalog<TRPCCatalog> {
+  private parseMessage(message: unknown):
+    | {
+        success: true;
+        data: MessagesFromSocketCatalog<TRPCCatalog>;
+      }
+    | {
+        success: false;
+        reason?: string;
+      } {
     const parsedMessage = messageSchema.safeParse(message);
 
     if (!parsedMessage.success) {
-      throw new Error(`Failed to parse message: ${JSON.stringify(parsedMessage.error)}`);
+      return {
+        success: false,
+        reason: `Failed to parse message: ${fromZodError(parsedMessage.error).toString()}`,
+      };
     }
 
     const schema = this.#schema[parsedMessage.data.type]["message"];
 
     if (!schema) {
-      throw new Error(`Unknown message type: ${parsedMessage.data.type}`);
+      return {
+        success: false,
+        reason: `Unknown message type: ${parsedMessage.data.type}`,
+      };
     }
 
     const messageWithVersion = {
@@ -141,14 +169,18 @@ export class ZodSocketMessageHandler<TRPCCatalog extends ZodSocketMessageCatalog
         payload: messageWithVersion,
       });
 
-      throw parsedPayload.error instanceof ZodError
-        ? fromZodError(parsedPayload.error)
-        : parsedPayload.error;
+      return {
+        success: false,
+        reason: fromZodError(parsedPayload.error).toString(),
+      };
     }
 
     return {
-      type: parsedMessage.data.type,
-      payload: parsedPayload.data,
+      success: true,
+      data: {
+        type: parsedMessage.data.type,
+        payload: parsedPayload.data,
+      },
     };
   }
 
@@ -202,6 +234,7 @@ export class ZodSocketMessageHandler<TRPCCatalog extends ZodSocketMessageCatalog
 export type ZodSocketMessageSenderOptions<TMessageCatalog extends ZodSocketMessageCatalogSchema> = {
   schema: TMessageCatalog;
   socket: ZodSocket<any, TMessageCatalog>;
+  logger?: StructuredLogger;
 };
 
 export type GetSocketMessagesWithCallback<TMessageCatalog extends ZodSocketMessageCatalogSchema> = {
@@ -221,10 +254,12 @@ export type GetSocketMessagesWithoutCallback<
 export class ZodSocketMessageSender<TMessageCatalog extends ZodSocketMessageCatalogSchema> {
   #schema: TMessageCatalog;
   #socket: ZodSocket<any, TMessageCatalog>;
+  #logger: StructuredLogger;
 
   constructor(options: ZodSocketMessageSenderOptions<TMessageCatalog>) {
     this.#schema = options.schema;
     this.#socket = options.socket;
+    this.#logger = options.logger ?? new SimpleStructuredLogger("zod-socket-sender", LogLevel.info);
   }
 
   public send<K extends GetSocketMessagesWithoutCallback<TMessageCatalog>>(
@@ -240,7 +275,10 @@ export class ZodSocketMessageSender<TMessageCatalog extends ZodSocketMessageCata
     const parsedPayload = schema.safeParse(payload);
 
     if (!parsedPayload.success) {
-      throw new Error(`Failed to parse message payload: ${JSON.stringify(parsedPayload.error)}`);
+      this.#logger.error("Failed to parse message payload, will not send", {
+        error: parsedPayload.error,
+      });
+      return;
     }
 
     // @ts-expect-error
@@ -275,10 +313,20 @@ export class ZodSocketMessageSender<TMessageCatalog extends ZodSocketMessageCata
 export type ZodSocket<
   TListenEvents extends ZodSocketMessageCatalogSchema,
   TEmitEvents extends ZodSocketMessageCatalogSchema,
-> = Socket<
-  ZodMessageCatalogToSocketIoEvents<TListenEvents>,
-  ZodMessageCatalogToSocketIoEvents<TEmitEvents>
->;
+> = Omit<
+  Socket<
+    ZodMessageCatalogToSocketIoEvents<TListenEvents>,
+    ZodMessageCatalogToSocketIoEvents<TEmitEvents>
+  >,
+  "timeout"
+> & {
+  timeout: (
+    timeout: number
+  ) => Socket<
+    ZodMessageCatalogToSocketIoEvents<TListenEvents>,
+    ZodMessageCatalogToSocketIoEvents<TEmitEvents>
+  >;
+};
 
 interface ZodSocketConnectionOptions<
   TClientMessages extends ZodSocketMessageCatalogSchema,
@@ -295,6 +343,7 @@ interface ZodSocketConnectionOptions<
   };
   handlers?: ZodSocketMessageHandlers<TServerMessages>;
   authToken?: string;
+  ioOptions?: Partial<ManagerOptions & SocketOptions>;
   onConnection?: (
     socket: ZodSocket<TServerMessages, TClientMessages>,
     handler: ZodSocketMessageHandler<TServerMessages>,
@@ -340,6 +389,7 @@ export class ZodSocketConnection<
       extraHeaders: opts.extraHeaders,
       reconnectionDelay: 500,
       reconnectionDelayMax: 1000,
+      ...opts.ioOptions,
     });
 
     this.#logger = logger.child({
@@ -355,6 +405,7 @@ export class ZodSocketConnection<
     this.#sender = new ZodSocketMessageSender({
       schema: opts.clientMessages,
       socket: this.socket,
+      logger: this.#logger,
     });
 
     this.socket.on("connect_error", async (error) => {
@@ -398,7 +449,3 @@ export class ZodSocketConnection<
     return this.#sender.sendWithAck.bind(this.#sender);
   }
 }
-
-function createLogger(prefix: string) {
-  return (...args: any[]) => console.log(prefix, ...args);
-}

pnpm-lock.yaml

@@ -92,9 +92,6 @@ importers:
       socket.io:
         specifier: 4.7.4
         version: 4.7.4
-      socket.io-client:
-        specifier: 4.7.4
-        version: 4.7.4
     devDependencies:
       '@types/node':
         specifier: ^18
@@ -123,9 +120,6 @@ importers:
       execa:
         specifier: ^8.0.1
         version: 8.0.1
-      socket.io-client:
-        specifier: ^4.7.4
-        version: 4.7.4
     devDependencies:
       '@types/node':
         specifier: ^18.19.8
@@ -157,9 +151,6 @@ importers:
       p-queue:
         specifier: ^8.0.1
         version: 8.0.1
-      socket.io-client:
-        specifier: ^4.7.4
-        version: 4.7.4
     devDependencies:
       dotenv:
         specifier: ^16.4.2
@@ -1624,9 +1615,6 @@ importers:
       simple-git:
         specifier: ^3.19.0
         version: 3.19.0
-      socket.io-client:
-        specifier: ^4.7.4
-        version: 4.7.4
       source-map-support:
         specifier: ^0.5.21
         version: 0.5.21
@@ -1758,8 +1746,8 @@ importers:
         specifier: ^3.27.3
         version: 3.27.3
       socket.io-client:
-        specifier: 4.7.4
-        version: 4.7.4
+        specifier: 4.7.5
+        version: 4.7.5
       superjson:
         specifier: ^2.2.1
         version: 2.2.1
@@ -1828,9 +1816,6 @@ importers:
       '@types/node':
         specifier: '18'
         version: 18.17.1
-      socket.io-client:
-        specifier: ^4.7.4
-        version: 4.7.4
       typescript:
         specifier: ^5.3.0
         version: 5.3.3
@@ -20459,6 +20444,7 @@ packages:
       - bufferutil
       - supports-color
       - utf-8-validate
+    dev: false
 
   /engine.io-parser@5.2.2(patch_hash=e6nctogrhpxoivwiwy37ersfu4):
     resolution: {integrity: sha512-RcyUFKA93/CXH20l4SoVvzZfrSDMOTUS3bWVpTt2FuFP+XYrL8i8oonHP7WInRyVHXh0n/ORtoeiE1os+8qkSw==}
@@ -32246,8 +32232,8 @@ packages:
       - utf-8-validate
     dev: false
 
-  /socket.io-client@4.7.4:
-    resolution: {integrity: sha512-wh+OkeF0rAVCrABWQBaEjLfb7DVPotMbu0cgWgyR0v6eA4EoVnAwcIeIbcdTE3GT/H3kbdLl7OoH2+asoDRIIg==}
+  /socket.io-client@4.7.5:
+    resolution: {integrity: sha512-sJ/tqHOCe7Z50JCBCXrsY3I2k03iOiUe+tj1OmKeD2lXPiGH/RUCdTZFoqVyN7l1MnpIzPrGtLcijffmeouNlQ==}
     engines: {node: '>=10.0.0'}
     dependencies:
       '@socket.io/component-emitter': 3.1.0
@@ -32258,6 +32244,7 @@ packages:
       - bufferutil
       - supports-color
       - utf-8-validate
+    dev: false
 
   /socket.io-parser@4.2.4:
     resolution: {integrity: sha512-/GbIKmo8ioc+NIWIhwdecY0ge+qVBSMdgxGygevmdHj24bsfgtCmcUUcQ5ZzcylGFHsN3k4HB4Cgkl96KVnuew==}
@@ -36393,6 +36380,7 @@ packages:
   /xmlhttprequest-ssl@2.0.0:
     resolution: {integrity: sha512-QKxVRxiRACQcVuQEYFsI1hhkrMlrXHPegbbd1yn9UHOmRxY+si12nQYzri3vbzt8VdTTRviqcKxcyllFas5z2A==}
     engines: {node: '>=0.4.0'}
+    dev: false
 
   /xtend@4.0.2:
     resolution: {integrity: sha512-LKYU1iAXJXUgAXn9URjiu+MWhyUXHsvfp7mcuYm9dSUKK0/CjtrUwFAxD82/mCWbtLsGjFIad0wIsod4zrTAEQ==}