More progress

Author: cte

packages/cloud/src/bridge/BridgeOrchestrator.ts

@@ -1,6 +1,8 @@
 import crypto from "crypto"
 import os from "os"
 
+import type { Socket } from "socket.io-client"
+
 import {
 	type TaskProviderLike,
 	type CloudUserInfo,
@@ -91,7 +93,7 @@ export class BridgeOrchestrator {
 				await options.provider.getTelemetryProperties()
 
 				BridgeOrchestrator.instance = new BridgeOrchestrator(options)
-				await BridgeOrchestrator.instance.connect()
+				await BridgeOrchestrator.instance.socketTransport.connect()
 			} catch (error) {
 				console.error(
 					`[BridgeOrchestrator#connectOrDisconnect] connect() failed: ${error instanceof Error ? error.message : String(error)}`,
@@ -103,7 +105,7 @@ export class BridgeOrchestrator {
 			if (connectionState === ConnectionState.FAILED || connectionState === ConnectionState.DISCONNECTED) {
 				console.log(`[BridgeOrchestrator#connectOrDisconnect] Re-connecting... (state: ${connectionState})`)
 
-				instance.reconnect().catch((error) => {
+				instance.socketTransport.reconnect().catch((error) => {
 					console.error(
 						`[BridgeOrchestrator#connectOrDisconnect] reconnect() failed: ${error instanceof Error ? error.message : String(error)}`,
 					)
@@ -124,7 +126,7 @@ export class BridgeOrchestrator {
 
 			try {
 				console.log(`[BridgeOrchestrator#connectOrDisconnect] Disconnecting... (state: ${connectionState})`)
-				await instance.disconnect()
+				await instance.socketTransport.disconnect()
 			} catch (error) {
 				console.error(
 					`[BridgeOrchestrator#connectOrDisconnect] disconnect() failed: ${error instanceof Error ? error.message : String(error)}`,
@@ -161,9 +163,18 @@ export class BridgeOrchestrator {
 				reconnectionDelay: this.RECONNECT_DELAY,
 				reconnectionDelayMax: this.RECONNECT_DELAY_MAX,
 			},
-			onConnect: () => this.handleConnect(),
-			onDisconnect: () => this.handleDisconnect(),
-			onReconnect: () => this.handleReconnect(),
+			onConnect: async (socket: Socket) => {
+				this.setupSocketListeners()
+				await this.extensionChannel.onConnect(socket)
+			},
+			onDisconnect: async () => {
+				await this.extensionChannel.onDisconnect()
+				await this.extensionChannel.cleanup(this.socketTransport.getSocket())
+			},
+			onReconnect: async (socket: Socket) => {
+				this.setupSocketListeners()
+				await this.extensionChannel.onReconnect(socket)
+			},
 		})
 
 		this.extensionChannel = new ExtensionChannel({
@@ -193,54 +204,7 @@ export class BridgeOrchestrator {
 				`[BridgeOrchestrator] on(${ExtensionSocketEvents.RELAYED_COMMAND}) -> ${message.type} for ${message.instanceId}`,
 			)
 
-			this.extensionChannel?.handleCommand(message)
+			this.extensionChannel.handleCommand(message)
 		})
 	}
-
-	private async handleConnect() {
-		const socket = this.socketTransport.getSocket()
-
-		if (!socket) {
-			console.error("[BridgeOrchestrator#handleConnect] Socket not available")
-			return
-		}
-
-		await this.extensionChannel.onConnect(socket)
-	}
-
-	private handleDisconnect() {
-		this.extensionChannel.onDisconnect()
-	}
-
-	private async handleReconnect() {
-		const socket = this.socketTransport.getSocket()
-
-		if (!socket) {
-			console.error("[BridgeOrchestrator] Socket not available after reconnect")
-			return
-		}
-
-		// Re-setup socket listeners to ensure they're properly configured
-		// after automatic reconnection (Socket.IO's built-in reconnection)
-		// The socket.off() calls in setupSocketListeners prevent duplicates
-		this.setupSocketListeners()
-
-		await this.extensionChannel.onReconnect(socket)
-	}
-
-	private async connect(): Promise<void> {
-		await this.socketTransport.connect()
-		this.setupSocketListeners()
-	}
-
-	private async disconnect(): Promise<void> {
-		await this.extensionChannel.cleanup(this.socketTransport.getSocket())
-		await this.socketTransport.disconnect()
-		BridgeOrchestrator.instance = null
-	}
-
-	private async reconnect(): Promise<void> {
-		await this.socketTransport.reconnect()
-		this.setupSocketListeners()
-	}
 }

packages/cloud/src/bridge/ExtensionChannel.ts

@@ -75,7 +75,7 @@ export class ExtensionChannel {
 		this.startHeartbeat(socket)
 	}
 
-	public onDisconnect(): void {
+	public async onDisconnect(): Promise<void> {
 		this.socket = null
 		this.stopHeartbeat()
 	}

packages/cloud/src/bridge/SocketTransport.ts

@@ -5,9 +5,9 @@ import { ConnectionState, type RetryConfig } from "@roo-code/types"
 export interface SocketTransportOptions {
 	url: string
 	socketOptions: Partial<ManagerOptions & SocketOptions>
-	onConnect?: () => void | Promise<void>
-	onDisconnect?: (reason: string) => void
-	onReconnect?: () => void | Promise<void>
+	onConnect?: (socket: Socket) => Promise<void>
+	onDisconnect?: (reason: string) => Promise<void>
+	onReconnect?: (socket: Socket) => Promise<void>
 	logger?: {
 		log: (message: string, ...args: unknown[]) => void
 		error: (message: string, ...args: unknown[]) => void
@@ -126,12 +126,12 @@ export class SocketTransport {
 				this.connectionState = ConnectionState.CONNECTED
 
 				if (this.isPreviouslyConnected) {
-					if (this.options.onReconnect) {
-						await this.options.onReconnect()
+					if (this.options.onReconnect && this.socket) {
+						await this.options.onReconnect(this.socket)
 					}
 				} else {
-					if (this.options.onConnect) {
-						await this.options.onConnect()
+					if (this.options.onConnect && this.socket) {
+						await this.options.onConnect(this.socket)
 					}
 				}
 
@@ -195,8 +195,8 @@ export class SocketTransport {
 				console.log(`[SocketTransport#_connect] on(reconnect) - ${attempt}`)
 				this.connectionState = ConnectionState.CONNECTED
 
-				if (this.options.onReconnect) {
-					this.options.onReconnect()
+				if (this.options.onReconnect && this.socket) {
+					this.options.onReconnect(this.socket)
 				}
 			})