Implement deferred task subscriptions (#7517)

Author: cte

packages/cloud/src/bridge/BridgeOrchestrator.ts

@@ -31,6 +31,8 @@ export interface BridgeOrchestratorOptions {
 export class BridgeOrchestrator {
 	private static instance: BridgeOrchestrator | null = null
 
+	private static pendingTask: TaskLike | null = null
+
 	// Core
 	private readonly userId: string
 	private readonly socketBridgeUrl: string
@@ -116,6 +118,22 @@ export class BridgeOrchestrator {
 		}
 	}
 
+	/**
+	 * @TODO: What if subtasks also get spawned? We'd probably want deferred
+	 * subscriptions for those too.
+	 */
+	public static async subscribeToTask(task: TaskLike): Promise<void> {
+		const instance = BridgeOrchestrator.instance
+
+		if (instance && instance.socketTransport.isConnected()) {
+			console.log(`[BridgeOrchestrator#subscribeToTask] Subscribing to task ${task.taskId}`)
+			await instance.subscribeToTask(task)
+		} else {
+			console.log(`[BridgeOrchestrator#subscribeToTask] Deferring subscription for task ${task.taskId}`)
+			BridgeOrchestrator.pendingTask = task
+		}
+	}
+
 	private constructor(options: BridgeOrchestratorOptions) {
 		this.userId = options.userId
 		this.socketBridgeUrl = options.socketBridgeUrl
@@ -180,12 +198,27 @@ export class BridgeOrchestrator {
 		const socket = this.socketTransport.getSocket()
 
 		if (!socket) {
-			console.error("[BridgeOrchestrator] Socket not available after connect")
+			console.error("[BridgeOrchestrator#handleConnect] Socket not available")
 			return
 		}
 
 		await this.extensionChannel.onConnect(socket)
 		await this.taskChannel.onConnect(socket)
+
+		if (BridgeOrchestrator.pendingTask) {
+			console.log(
+				`[BridgeOrchestrator#handleConnect] Subscribing to task ${BridgeOrchestrator.pendingTask.taskId}`,
+			)
+
+			try {
+				await this.subscribeToTask(BridgeOrchestrator.pendingTask)
+				BridgeOrchestrator.pendingTask = null
+			} catch (error) {
+				console.error(
+					`[BridgeOrchestrator#handleConnect] subscribeToTask() failed: ${error instanceof Error ? error.message : String(error)}`,
+				)
+			}
+		}
 	}
 
 	private handleDisconnect() {
@@ -261,6 +294,7 @@ export class BridgeOrchestrator {
 		await this.taskChannel.cleanup(this.socketTransport.getSocket())
 		await this.socketTransport.disconnect()
 		BridgeOrchestrator.instance = null
+		BridgeOrchestrator.pendingTask = null
 	}
 
 	public async reconnect(): Promise<void> {

packages/cloud/src/bridge/TaskChannel.ts

@@ -163,25 +163,27 @@ export class TaskChannel extends BaseChannel<
 	public async unsubscribeFromTask(taskId: string, _socket: Socket): Promise<void> {
 		const task = this.subscribedTasks.get(taskId)
 
+		if (!task) {
+			return
+		}
+
 		await this.publish(TaskSocketEvents.LEAVE, { taskId }, (response: LeaveResponse) => {
 			if (response.success) {
-				console.log(`[TaskChannel#unsubscribeFromTask] unsubscribed from ${taskId}`, response)
+				console.log(`[TaskChannel#unsubscribeFromTask] unsubscribed from ${taskId}`)
 			} else {
 				console.error(`[TaskChannel#unsubscribeFromTask] failed to unsubscribe from ${taskId}`)
 			}
 
 			// If we failed to unsubscribe then something is probably wrong and
 			// we should still discard this task from `subscribedTasks`.
-			if (task) {
-				this.removeTaskListeners(task)
-				this.subscribedTasks.delete(taskId)
-			}
+			this.removeTaskListeners(task)
+			this.subscribedTasks.delete(taskId)
 		})
 	}
 
 	private setupTaskListeners(task: TaskLike): void {
 		if (this.taskListeners.has(task.taskId)) {
-			console.warn("[TaskChannel] Listeners already exist for task, removing old listeners:", task.taskId)
+			console.warn(`[TaskChannel] Listeners already exist for task, removing old listeners for ${task.taskId}`)
 			this.removeTaskListeners(task)
 		}
 

packages/cloud/src/bridge/__tests__/TaskChannel.test.ts

@@ -299,8 +299,7 @@ describe("TaskChannel", () => {
 
 			// Verify warning was logged
 			expect(warnSpy).toHaveBeenCalledWith(
-				"[TaskChannel] Listeners already exist for task, removing old listeners:",
-				taskId,
+				`[TaskChannel] Listeners already exist for task, removing old listeners for ${taskId}`,
 			)
 
 			// Verify only one set of listeners exists

packages/types/npm/package.metadata.json

@@ -1,6 +1,6 @@
 {
 	"name": "@roo-code/types",
-	"version": "1.64.0",
+	"version": "1.65.0",
 	"description": "TypeScript type definitions for Roo Code.",
 	"publishConfig": {
 		"access": "public",

src/core/task/Task.ts

@@ -258,7 +258,6 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 
 	// Task Bridge
 	enableBridge: boolean
-	bridge: BridgeOrchestrator | null = null
 
 	// Streaming
 	isWaitingForFirstChunk = false
@@ -1086,14 +1085,10 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 	private async startTask(task?: string, images?: string[]): Promise<void> {
 		if (this.enableBridge) {
 			try {
-				this.bridge = this.bridge || BridgeOrchestrator.getInstance()
-
-				if (this.bridge) {
-					await this.bridge.subscribeToTask(this)
-				}
+				await BridgeOrchestrator.subscribeToTask(this)
 			} catch (error) {
 				console.error(
-					`[Task#startTask] subscribeToTask failed - ${error instanceof Error ? error.message : String(error)}`,
+					`[Task#startTask] BridgeOrchestrator.subscribeToTask() failed: ${error instanceof Error ? error.message : String(error)}`,
 				)
 			}
 		}
@@ -1158,14 +1153,10 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 	private async resumeTaskFromHistory() {
 		if (this.enableBridge) {
 			try {
-				this.bridge = this.bridge || BridgeOrchestrator.getInstance()
-
-				if (this.bridge) {
-					await this.bridge.subscribeToTask(this)
-				}
+				await BridgeOrchestrator.subscribeToTask(this)
 			} catch (error) {
 				console.error(
-					`[Task#resumeTaskFromHistory] subscribeToTask failed - ${error instanceof Error ? error.message : String(error)}`,
+					`[Task#resumeTaskFromHistory] BridgeOrchestrator.subscribeToTask() failed: ${error instanceof Error ? error.message : String(error)}`,
 				)
 			}
 		}
@@ -1419,10 +1410,9 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 	}
 
 	public dispose(): void {
-		// Disposing task
-		console.log(`[Task] disposing task ${this.taskId}.${this.instanceId}`)
+		console.log(`[Task#dispose] disposing task ${this.taskId}.${this.instanceId}`)
 
-		// Remove all event listeners to prevent memory leaks
+		// Remove all event listeners to prevent memory leaks.
 		try {
 			this.removeAllListeners()
 		} catch (error) {
@@ -1442,13 +1432,14 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 			this.pauseInterval = undefined
 		}
 
-		// Unsubscribe from TaskBridge service.
-		if (this.bridge) {
-			this.bridge
-				.unsubscribeFromTask(this.taskId)
-				.catch((error: unknown) => console.error("Error unsubscribing from task bridge:", error))
-
-			this.bridge = null
+		if (this.enableBridge) {
+			BridgeOrchestrator.getInstance()
+				?.unsubscribeFromTask(this.taskId)
+				.catch((error) =>
+					console.error(
+						`[Task#dispose] BridgeOrchestrator#unsubscribeFromTask() failed: ${error instanceof Error ? error.message : String(error)}`,
+					),
+				)
 		}
 
 		// Release any terminals associated with this task.

src/core/webview/ClineProvider.ts

@@ -167,6 +167,8 @@ export class ClineProvider
 
 		this.marketplaceManager = new MarketplaceManager(this.context, this.customModesManager)
 
+		// Forward <most> task events to the provider.
+		// We do something fairly similar for the IPC-based API.
 		this.taskCreationCallback = (instance: Task) => {
 			this.emit(RooCodeEventName.TaskCreated, instance)
 
@@ -348,18 +350,18 @@ export class ClineProvider
 		let task = this.clineStack.pop()
 
 		if (task) {
+			task.emit(RooCodeEventName.TaskUnfocused)
+
 			try {
 				// Abort the running task and set isAbandoned to true so
 				// all running promises will exit as well.
 				await task.abortTask(true)
 			} catch (e) {
 				this.log(
-					`[removeClineFromStack] encountered error while aborting task ${task.taskId}.${task.instanceId}: ${e.message}`,
+					`[ClineProvider#removeClineFromStack] abortTask() failed ${task.taskId}.${task.instanceId}: ${e.message}`,
 				)
 			}
 
-			task.emit(RooCodeEventName.TaskUnfocused)
-
 			// Remove event listeners before clearing the reference.
 			const cleanupFunctions = this.taskEventListeners.get(task)
 
@@ -407,12 +409,6 @@ export class ClineProvider
 		await this.getCurrentTask()?.resumePausedTask(lastMessage)
 	}
 
-	// Clear the current task without treating it as a subtask.
-	// This is used when the user cancels a task that is not a subtask.
-	async clearTask() {
-		await this.removeClineFromStack()
-	}
-
 	resumeTask(taskId: string): void {
 		// Use the existing showTaskWithId method which handles both current and historical tasks
 		this.showTaskWithId(taskId).catch((error) => {
@@ -1365,6 +1361,16 @@ export class ClineProvider
 		await this.createTaskWithHistoryItem({ ...historyItem, rootTask, parentTask, preservedFCOState })
 	}
 
+	// Clear the current task without treating it as a subtask.
+	// This is used when the user cancels a task that is not a subtask.
+	async clearTask() {
+		if (this.clineStack.length > 0) {
+			const task = this.clineStack[this.clineStack.length - 1]
+			console.log(`[clearTask] clearing task ${task.taskId}.${task.instanceId}`)
+			await this.removeClineFromStack()
+		}
+	}
+
 	async updateCustomInstructions(instructions?: string) {
 		// User may be clearing the field.
 		await this.updateGlobalState("customInstructions", instructions || undefined)
@@ -1643,6 +1649,7 @@ export class ClineProvider
 			})
 		} catch (error) {
 			console.error("Failed to fetch marketplace data:", error)
+
 			// Send empty data on error to prevent UI from hanging
 			this.postMessageToWebview({
 				type: "marketplaceData",
@@ -2272,24 +2279,23 @@ export class ClineProvider
 		if (bridge) {
 			const currentTask = this.getCurrentTask()
 
-			if (currentTask && !currentTask.bridge) {
+			if (currentTask && !currentTask.enableBridge) {
 				try {
-					currentTask.bridge = bridge
-					await currentTask.bridge.subscribeToTask(currentTask)
+					currentTask.enableBridge = true
+					await BridgeOrchestrator.subscribeToTask(currentTask)
 				} catch (error) {
-					const message = `[ClineProvider#remoteControlEnabled] subscribeToTask failed - ${error instanceof Error ? error.message : String(error)}`
+					const message = `[ClineProvider#remoteControlEnabled] BridgeOrchestrator.subscribeToTask() failed: ${error instanceof Error ? error.message : String(error)}`
 					this.log(message)
 					console.error(message)
 				}
 			}
 		} else {
 			for (const task of this.clineStack) {
-				if (task.bridge) {
+				if (task.enableBridge) {
 					try {
-						await task.bridge.unsubscribeFromTask(task.taskId)
-						task.bridge = null
+						await BridgeOrchestrator.getInstance()?.unsubscribeFromTask(task.taskId)
 					} catch (error) {
-						const message = `[ClineProvider#remoteControlEnabled] unsubscribeFromTask failed - ${error instanceof Error ? error.message : String(error)}`
+						const message = `[ClineProvider#remoteControlEnabled] BridgeOrchestrator#unsubscribeFromTask() failed: ${error instanceof Error ? error.message : String(error)}`
 						this.log(message)
 						console.error(message)
 					}