Merge pull request #3 from BasicBlock/pn/upgrade

potential fix for socket disconnect

Author: pnispel

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

@@ -468,6 +468,18 @@ export class ManagedRunController {
         runId: this.runFriendlyId,
         message: "Socket connected to supervisor",
       });
+
+      if (this.runFriendlyId && this.snapshotFriendlyId) {
+        this.sendDebugLog({
+          runId: this.runFriendlyId,
+          message: "Socket connected - re-subscribing to run notifications",
+          properties: {
+            snapshotFriendlyId: this.snapshotFriendlyId,
+          },
+        });
+
+        this.subscribeToRunNotifications(this.runFriendlyId, this.snapshotFriendlyId);
+      }
     });
 
     socket.on("connect_error", (error) => {
@@ -496,6 +508,22 @@ export class ManagedRunController {
     });
 
     socket.on("disconnect", async (reason, description) => {
+      const reconnectSocket = (trigger: string) => {
+        this.sendDebugLog({
+          runId: this.runFriendlyId,
+          message: "Reconnecting supervisor socket",
+          properties: { trigger, reason },
+        });
+
+        socket.removeAllListeners();
+        socket.disconnect();
+        this.socket = this.createSupervisorSocket();
+
+        if (this.runFriendlyId && this.snapshotFriendlyId) {
+          this.subscribeToRunNotifications(this.runFriendlyId, this.snapshotFriendlyId);
+        }
+      };
+
       const parseDescription = ():
         | {
             description: string;
@@ -539,41 +567,31 @@ export class ManagedRunController {
           properties: { reason, ...parseDescription(), currentEnv, newEnv },
         });
 
+        const shouldForceReconnect = reason === "ping timeout";
+
         if (!result) {
+          if (shouldForceReconnect) {
+            reconnectSocket("disconnect override fetch failed");
+          }
           return;
         }
 
-        // If runner ID changed, we detected a restore
-        if (result.runnerIdChanged) {
+        if (result.runnerIdChanged || result.supervisorChanged) {
           this.sendDebugLog({
             runId: this.runFriendlyId,
-            message: "Runner ID changed - restore detected",
+            message: "Restore-related env override detected",
             properties: {
+              runnerIdChanged: result.runnerIdChanged,
               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();
+          reconnectSocket("restore-related env override");
+          return;
+        }
 
-          // Re-subscribe to notifications if we have an active execution
-          if (this.runFriendlyId && this.snapshotFriendlyId) {
-            this.subscribeToRunNotifications(this.runFriendlyId, this.snapshotFriendlyId);
-          }
+        if (shouldForceReconnect) {
+          reconnectSocket("ping timeout");
         }
 
         return;