PR feedback

packages/cloud/src/bridge/ExtensionChannel.ts

@@ -2,6 +2,7 @@ import type { Socket } from "socket.io-client"
 
 import {
 	type TaskProviderLike,
+	type TaskProviderEvents,
 	type ExtensionInstance,
 	type ExtensionBridgeCommand,
 	type ExtensionBridgeEvent,
@@ -28,6 +29,7 @@ export class ExtensionChannel extends BaseChannel<
 	private provider: TaskProviderLike
 	private extensionInstance: ExtensionInstance
 	private heartbeatInterval: NodeJS.Timeout | null = null
+	private eventListeners: Map<RooCodeEventName, (...args: unknown[]) => void> = new Map()
 
 	constructor(instanceId: string, userId: string, provider: TaskProviderLike) {
 		super(instanceId)
@@ -115,6 +117,7 @@ export class ExtensionChannel extends BaseChannel<
 
 	protected async handleCleanup(socket: Socket): Promise<void> {
 		this.stopHeartbeat()
+		this.cleanupListeners()
 		await this.unregisterInstance(socket)
 	}
 
@@ -168,17 +171,28 @@ export class ExtensionChannel extends BaseChannel<
 			{ from: RooCodeEventName.TaskIdle, to: ExtensionBridgeEventName.TaskIdle },
 		] as const
 
-		const addListener =
-			(type: ExtensionBridgeEventName) =>
-			(..._args: unknown[]) => {
+		eventMapping.forEach(({ from, to }) => {
+			// Create and store the listener function for cleanup/
+			const listener = (..._args: unknown[]) => {
 				this.publish(ExtensionSocketEvents.EVENT, {
-					type,
+					type: to,
 					instance: this.updateInstance(),
 					timestamp: Date.now(),
 				})
 			}
 
-		eventMapping.forEach(({ from, to }) => this.provider.on(from, addListener(to)))
+			this.eventListeners.set(from, listener)
+			this.provider.on(from, listener)
+		})
+	}
+
+	private cleanupListeners(): void {
+		this.eventListeners.forEach((listener, eventName) => {
+			// Cast is safe because we only store valid event names from eventMapping.
+			this.provider.off(eventName as keyof TaskProviderEvents, listener)
+		})
+
+		this.eventListeners.clear()
 	}
 
 	private updateInstance(): ExtensionInstance {

packages/cloud/src/bridge/TaskChannel.ts

@@ -20,10 +20,10 @@ type TaskEventListener = {
 	[K in keyof TaskEvents]: (...args: TaskEvents[K]) => void | Promise<void>
 }[keyof TaskEvents]
 
-const TASK_EVENT_MAPPING: Record<TaskBridgeEventName, keyof TaskEvents> = {
-	[TaskBridgeEventName.Message]: RooCodeEventName.Message,
-	[TaskBridgeEventName.TaskModeSwitched]: RooCodeEventName.TaskModeSwitched,
-	[TaskBridgeEventName.TaskInteractive]: RooCodeEventName.TaskInteractive,
+type TaskEventMapping = {
+	from: keyof TaskEvents
+	to: TaskBridgeEventName
+	createPayload: (task: TaskLike, ...args: any[]) => any // eslint-disable-line @typescript-eslint/no-explicit-any
 }
 
 /**
@@ -39,6 +39,36 @@ export class TaskChannel extends BaseChannel<
 	private pendingTasks: Map<string, TaskLike> = new Map()
 	private taskListeners: Map<string, Map<TaskBridgeEventName, TaskEventListener>> = new Map()
 
+	private readonly eventMapping: readonly TaskEventMapping[] = [
+		{
+			from: RooCodeEventName.Message,
+			to: TaskBridgeEventName.Message,
+			createPayload: (task: TaskLike, data: { action: string; message: ClineMessage }) => ({
+				type: TaskBridgeEventName.Message,
+				taskId: task.taskId,
+				action: data.action,
+				message: data.message,
+			}),
+		},
+		{
+			from: RooCodeEventName.TaskModeSwitched,
+			to: TaskBridgeEventName.TaskModeSwitched,
+			createPayload: (task: TaskLike, mode: string) => ({
+				type: TaskBridgeEventName.TaskModeSwitched,
+				taskId: task.taskId,
+				mode,
+			}),
+		},
+		{
+			from: RooCodeEventName.TaskInteractive,
+			to: TaskBridgeEventName.TaskInteractive,
+			createPayload: (task: TaskLike, _taskId: string) => ({
+				type: TaskBridgeEventName.TaskInteractive,
+				taskId: task.taskId,
+			}),
+		},
+	] as const
+
 	constructor(instanceId: string) {
 		super(instanceId)
 	}
@@ -157,35 +187,16 @@ export class TaskChannel extends BaseChannel<
 
 		const listeners = new Map<TaskBridgeEventName, TaskEventListener>()
 
-		const onMessage = ({ action, message }: { action: string; message: ClineMessage }) => {
-			this.publish(TaskSocketEvents.EVENT, {
-				type: TaskBridgeEventName.Message,
-				taskId: task.taskId,
-				action,
-				message,
-			})
-		}
-		task.on(RooCodeEventName.Message, onMessage)
-		listeners.set(TaskBridgeEventName.Message, onMessage)
-
-		const onTaskModeSwitched = (mode: string) => {
-			this.publish(TaskSocketEvents.EVENT, {
-				type: TaskBridgeEventName.TaskModeSwitched,
-				taskId: task.taskId,
-				mode,
-			})
-		}
-		task.on(RooCodeEventName.TaskModeSwitched, onTaskModeSwitched)
-		listeners.set(TaskBridgeEventName.TaskModeSwitched, onTaskModeSwitched)
+		this.eventMapping.forEach(({ from, to, createPayload }) => {
+			const listener = (...args: unknown[]) => {
+				const payload = createPayload(task, ...args)
+				this.publish(TaskSocketEvents.EVENT, payload)
+			}
 
-		const onTaskInteractive = (_taskId: string) => {
-			this.publish(TaskSocketEvents.EVENT, {
-				type: TaskBridgeEventName.TaskInteractive,
-				taskId: task.taskId,
-			})
-		}
-		task.on(RooCodeEventName.TaskInteractive, onTaskInteractive)
-		listeners.set(TaskBridgeEventName.TaskInteractive, onTaskInteractive)
+			// eslint-disable-next-line @typescript-eslint/no-explicit-any
+			task.on(from, listener as any)
+			listeners.set(to, listener)
+		})
 
 		this.taskListeners.set(task.taskId, listeners)
 	}
@@ -197,14 +208,18 @@ export class TaskChannel extends BaseChannel<
 			return
 		}
 
-		listeners.forEach((listener, eventName) => {
-			try {
-				// eslint-disable-next-line @typescript-eslint/no-explicit-any
-				task.off(TASK_EVENT_MAPPING[eventName], listener as any)
-			} catch (error) {
-				console.error(
-					`[TaskChannel] task.off(${TASK_EVENT_MAPPING[eventName]}) failed for task ${task.taskId}: ${error instanceof Error ? error.message : String(error)}`,
-				)
+		this.eventMapping.forEach(({ from, to }) => {
+			const listener = listeners.get(to)
+			if (listener) {
+				try {
+					task.off(from, listener as any) // eslint-disable-line @typescript-eslint/no-explicit-any
+				} catch (error) {
+					console.error(
+						`[TaskChannel] task.off(${from}) failed for task ${task.taskId}: ${
+							error instanceof Error ? error.message : String(error)
+						}`,
+					)
+				}
 			}
 		})