diff --git a/packages/cli-v3/src/entryPoints/managed/controller.ts b/packages/cli-v3/src/entryPoints/managed/controller.ts index 29ade1628de..8a028948873 100644 --- a/packages/cli-v3/src/entryPoints/managed/controller.ts +++ b/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;