fix(runner): improve restore detection

nicktrn · nicktrn · commit aad78408b4ac · 2025-09-17T19:51:20.000+01:00
diff --git a/packages/cli-v3/src/entryPoints/managed/controller.ts b/packages/cli-v3/src/entryPoints/managed/controller.ts
@@ -461,19 +461,6 @@ export class ManagedRunController {
         runId: this.runFriendlyId,
         message: "Socket connected to supervisor",
       });
-
-      // This should handle the case where we reconnect after being restored
-      if (
-        this.runFriendlyId &&
-        this.snapshotFriendlyId &&
-        this.runFriendlyId !== this.env.TRIGGER_RUN_ID
-      ) {
-        this.sendDebugLog({
-          runId: this.runFriendlyId,
-          message: "Subscribing to notifications for in-progress run",
-        });
-        this.subscribeToRunNotifications(this.runFriendlyId, this.snapshotFriendlyId);
-      }
     });
 
     socket.on("connect_error", (error) => {
@@ -514,7 +501,7 @@ export class ManagedRunController {
           supervisorApiUrl: this.env.TRIGGER_SUPERVISOR_API_URL,
         };
 
-        await this.currentExecution.processEnvOverrides("socket disconnected", true);
+        const result = await this.currentExecution.processEnvOverrides("socket disconnected", true);
 
         const newEnv = {
           workerInstanceName: this.env.TRIGGER_WORKER_INSTANCE_NAME,
@@ -528,6 +515,43 @@ export class ManagedRunController {
           properties: { reason, ...parseDescription(), currentEnv, newEnv },
         });
 
+        if (!result) {
+          return;
+        }
+
+        // If runner ID changed, we detected a restore
+        if (result.runnerIdChanged) {
+          this.sendDebugLog({
+            runId: this.runFriendlyId,
+            message: "Runner ID changed - restore detected",
+            properties: {
+              supervisorChanged: result.supervisorChanged,
+            },
+          });
+
+          if (!result.supervisorChanged) {
+            return;
+          }
+
+          // Only reconnect WebSocket if supervisor URL actually changed
+          this.sendDebugLog({
+            runId: this.runFriendlyId,
+            message: "Supervisor URL changed - creating new socket connection",
+          });
+
+          // First disconnect the old socket to avoid conflicts
+          socket.removeAllListeners();
+          socket.disconnect();
+
+          // Create a new socket with the updated URL and headers
+          this.socket = this.createSupervisorSocket();
+
+          // Re-subscribe to notifications if we have an active execution
+          if (this.runFriendlyId && this.snapshotFriendlyId) {
+            this.subscribeToRunNotifications(this.runFriendlyId, this.snapshotFriendlyId);
+          }
+        }
+
         return;
       }
 
diff --git a/packages/cli-v3/src/entryPoints/managed/execution.ts b/packages/cli-v3/src/entryPoints/managed/execution.ts
@@ -873,7 +873,23 @@ export class RunExecution {
     );
 
     if (!continuationResult.success) {
-      throw new Error(continuationResult.error);
+      // Check if we need to refresh metadata due to connection error
+      if (continuationResult.isConnectionError) {
+        this.sendDebugLog("restore: connection error detected, refreshing metadata");
+        await this.processEnvOverrides("restore connection error");
+
+        // Retry the continuation after refreshing metadata
+        const retryResult = await this.httpClient.continueRunExecution(
+          this.runFriendlyId,
+          this.snapshotManager.snapshotId
+        );
+
+        if (!retryResult.success) {
+          throw new Error(retryResult.error);
+        }
+      } else {
+        throw new Error(continuationResult.error);
+      }
     }
 
     // Track restore count
@@ -899,11 +915,18 @@ export class RunExecution {
   public async processEnvOverrides(
     reason?: string,
     shouldPollForSnapshotChanges?: boolean
-  ): Promise<{ overrides: Metadata } | null> {
+  ): Promise<{
+    overrides: Metadata;
+    runnerIdChanged?: boolean;
+    supervisorChanged?: boolean;
+  } | null> {
     if (!this.metadataClient) {
       return null;
     }
 
+    const previousRunnerId = this.env.TRIGGER_RUNNER_ID;
+    const previousSupervisorUrl = this.env.TRIGGER_SUPERVISOR_API_URL;
+
     const [error, overrides] = await this.metadataClient.getEnvOverrides();
 
     if (error) {
@@ -931,6 +954,14 @@ export class RunExecution {
     // Override the env with the new values
     this.env.override(overrides);
 
+    // Check if runner ID changed
+    const newRunnerId = this.env.TRIGGER_RUNNER_ID;
+    const runnerIdChanged = previousRunnerId !== newRunnerId;
+
+    // Check if supervisor URL changed
+    const newSupervisorUrl = this.env.TRIGGER_SUPERVISOR_API_URL;
+    const supervisorChanged = previousSupervisorUrl !== newSupervisorUrl;
+
     // Update services with new values
     if (overrides.TRIGGER_SNAPSHOT_POLL_INTERVAL_SECONDS) {
       this.snapshotPoller?.updateInterval(this.env.TRIGGER_SNAPSHOT_POLL_INTERVAL_SECONDS * 1000);
@@ -954,6 +985,8 @@ export class RunExecution {
 
     return {
       overrides,
+      runnerIdChanged,
+      supervisorChanged,
     };
   }
 
@@ -977,6 +1010,12 @@ export class RunExecution {
 
     if (!response.success) {
       this.sendDebugLog("heartbeat: failed", { error: response.error });
+
+      // Check if we need to refresh metadata due to connection error
+      if (response.isConnectionError) {
+        this.sendDebugLog("heartbeat: connection error detected, refreshing metadata");
+        await this.processEnvOverrides("heartbeat connection error");
+      }
     }
 
     this.lastHeartbeat = new Date();
@@ -1192,6 +1231,14 @@ export class RunExecution {
         error: response.error,
       });
 
+      if (response.isConnectionError) {
+        // Log this separately to make it more visible
+        this.sendDebugLog(
+          "fetchAndProcessSnapshotChanges: connection error detected, refreshing metadata"
+        );
+      }
+
+      // Always trigger metadata refresh on snapshot fetch errors
       await this.processEnvOverrides("snapshots since error");
       return;
     }
diff --git a/packages/core/src/v3/runEngineWorker/workload/http.ts b/packages/core/src/v3/runEngineWorker/workload/http.ts
@@ -52,18 +52,58 @@ export class WorkloadHttpClient {
     });
   }
 
+  private isConnectionError(error: string): boolean {
+    const connectionErrors = [
+      "Connection error",
+      "ECONNREFUSED",
+      "ETIMEDOUT",
+      "ENOTFOUND",
+      "ECONNRESET",
+      "EHOSTUNREACH",
+      "ENETUNREACH",
+      "EPIPE",
+      "ECONNABORTED",
+    ];
+    return connectionErrors.some((errType) => error.includes(errType));
+  }
+
+  private async withConnectionErrorDetection<T>(
+    operation: () => Promise<{ success: true; data: T } | { success: false; error: string }>
+  ): Promise<
+    { success: true; data: T } | { success: false; error: string; isConnectionError?: boolean }
+  > {
+    const result = await operation();
+
+    if (result.success) {
+      return result;
+    }
+
+    // Check if this is a connection error
+    if (this.isConnectionError(result.error)) {
+      return {
+        ...result,
+        isConnectionError: true,
+      };
+    }
+
+    return result;
+  }
+
   async heartbeatRun(runId: string, snapshotId: string, body?: WorkloadHeartbeatRequestBody) {
-    return wrapZodFetch(
-      WorkloadHeartbeatResponseBody,
-      `${this.apiUrl}/api/v1/workload-actions/runs/${runId}/snapshots/${snapshotId}/heartbeat`,
-      {
-        method: "POST",
-        headers: {
-          ...this.defaultHeaders(),
-          "Content-Type": "application/json",
-        },
-        body: JSON.stringify(body ?? {}),
-      }
+    return this.withConnectionErrorDetection(() =>
+      wrapZodFetch(
+        WorkloadHeartbeatResponseBody,
+        `${this.apiUrl}/api/v1/workload-actions/runs/${runId}/snapshots/${snapshotId}/heartbeat`,
+        {
+          method: "POST",
+          headers: {
+            ...this.defaultHeaders(),
+            "Content-Type": "application/json",
+          },
+          body: JSON.stringify(body ?? {}),
+          signal: AbortSignal.timeout(10_000), // 10 second timeout
+        }
+      )
     );
   }
 
@@ -81,15 +121,17 @@ export class WorkloadHttpClient {
   }
 
   async continueRunExecution(runId: string, snapshotId: string) {
-    return wrapZodFetch(
-      WorkloadContinueRunExecutionResponseBody,
-      `${this.apiUrl}/api/v1/workload-actions/runs/${runId}/snapshots/${snapshotId}/continue`,
-      {
-        method: "GET",
-        headers: {
-          ...this.defaultHeaders(),
-        },
-      }
+    return this.withConnectionErrorDetection(() =>
+      wrapZodFetch(
+        WorkloadContinueRunExecutionResponseBody,
+        `${this.apiUrl}/api/v1/workload-actions/runs/${runId}/snapshots/${snapshotId}/continue`,
+        {
+          method: "GET",
+          headers: {
+            ...this.defaultHeaders(),
+          },
+        }
+      )
     );
   }
 
@@ -130,15 +172,18 @@ export class WorkloadHttpClient {
   }
 
   async getSnapshotsSince(runId: string, snapshotId: string) {
-    return wrapZodFetch(
-      WorkloadRunSnapshotsSinceResponseBody,
-      `${this.apiUrl}/api/v1/workload-actions/runs/${runId}/snapshots/since/${snapshotId}`,
-      {
-        method: "GET",
-        headers: {
-          ...this.defaultHeaders(),
-        },
-      }
+    return this.withConnectionErrorDetection(() =>
+      wrapZodFetch(
+        WorkloadRunSnapshotsSinceResponseBody,
+        `${this.apiUrl}/api/v1/workload-actions/runs/${runId}/snapshots/since/${snapshotId}`,
+        {
+          method: "GET",
+          headers: {
+            ...this.defaultHeaders(),
+          },
+          signal: AbortSignal.timeout(10_000), // 10 second timeout
+        }
+      )
     );
   }
 
