fix retry mechanics

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

@@ -103,7 +103,7 @@ export class ProdBackgroundWorker {
     await this._taskRunProcess?.cleanup(true, gracefulExitTimeoutElapsed);
   }
 
-  async killTaskRunProcess(flush = true, initialSignal: number | NodeJS.Signals = "SIGTERM") {
+  async #killTaskRunProcess(flush = true, initialSignal: number | NodeJS.Signals = "SIGTERM") {
     if (this._closed || !this._taskRunProcess) {
       return;
     }
@@ -250,7 +250,11 @@ export class ProdBackgroundWorker {
     );
 
     taskRunProcess.onExit.attach(({ pid }) => {
-      this._taskRunProcess = undefined;
+      // Only delete the task run process if the pid matches
+      if (this._taskRunProcess?.pid === pid) {
+        this._taskRunProcess = undefined;
+      }
+
       if (pid) {
         this._taskRunProcessesBeingKilled.delete(pid);
       }
@@ -467,6 +471,8 @@ export class ProdBackgroundWorker {
           code: TaskRunErrorCodes.TASK_EXECUTION_FAILED,
         },
       };
+    } finally {
+      await this.#killTaskRunProcess();
     }
   }
 

packages/cli-v3/src/workers/prod/entry-point.ts

@@ -314,19 +314,22 @@ class ProdWorker {
     this.executing = false;
     this.attemptFriendlyId = undefined;
 
-    // Every retry gets a fresh process
-    await this.#backgroundWorker.killTaskRunProcess();
-
     if (willCheckpointAndRestore) {
       this.waitForPostStart = true;
+
+      // We already flush after completion, so we don't need to do it here
+      this.#prepareForCheckpoint(false);
+
       this.#coordinatorSocket.socket.emit("READY_FOR_CHECKPOINT", { version: "v1" });
       return;
     }
   }
 
-  async #prepareForCheckpoint() {
-    // Flush before checkpointing so we don't flush the same spans again after restore
-    await this.#backgroundWorker.flushTelemetry();
+  async #prepareForCheckpoint(flush = true) {
+    if (flush) {
+      // Flush before checkpointing so we don't flush the same spans again after restore
+      await this.#backgroundWorker.flushTelemetry();
+    }
 
     // Kill the previous worker process to prevent large checkpoints
     await this.#backgroundWorker.forceKillOldTaskRunProcesses();