no-mistakes(review): await initialization in sendToExtension consistently

Author: taltas

apps/cli/src/agent/__tests__/extension-host.test.ts

@@ -294,35 +294,52 @@ describe("ExtensionHost", () => {
 	})
 
 	describe("sendToExtension", () => {
-		it("should throw error when extension not ready", () => {
+		it("should throw error when extension not ready", async () => {
 			const host = createTestHost()
 			const message: WebviewMessage = { type: "requestModes" }
 
-			expect(() => {
-				host.sendToExtension(message)
-			}).toThrow("You cannot send messages to the extension before it is ready")
+			await expect(host.sendToExtension(message)).rejects.toThrow(
+				"You cannot send messages to the extension before it is ready",
+			)
 		})
 
-		it("should emit webviewMessage event when webview is ready", () => {
+		it("should emit webviewMessage event when webview is ready", async () => {
 			const host = createTestHost()
 			const emitSpy = vi.spyOn(host, "emit")
 			const message: WebviewMessage = { type: "requestModes" }
 
 			host.markWebviewReady()
 			emitSpy.mockClear() // Clear the markWebviewReady calls
-			host.sendToExtension(message)
+			await host.sendToExtension(message)
 
 			expect(emitSpy).toHaveBeenCalledWith("webviewMessage", message)
 		})
 
-		it("should not throw when webview is ready", () => {
+		it("should not throw when webview is ready", async () => {
 			const host = createTestHost()
 
 			host.markWebviewReady()
 
-			expect(() => {
-				host.sendToExtension({ type: "requestModes" })
-			}).not.toThrow()
+			await expect(host.sendToExtension({ type: "requestModes" })).resolves.not.toThrow()
+		})
+
+		it("should await initialization before emitting message", async () => {
+			const host = createTestHost()
+			const emitSpy = vi.spyOn(host, "emit")
+			const message: WebviewMessage = { type: "requestModes" }
+
+			host.markWebviewReady()
+			emitSpy.mockClear()
+
+			// Send message immediately after marking ready
+			const promise = host.sendToExtension(message)
+
+			// Message should not be emitted yet if initialization is still pending
+			// In this test, initialization completes synchronously, so this is just
+			// ensuring the await happens
+			await promise
+
+			expect(emitSpy).toHaveBeenCalledWith("webviewMessage", message)
 		})
 	})
 
@@ -343,6 +360,49 @@ describe("ExtensionHost", () => {
 		})
 	})
 
+	describe("initialization ordering", () => {
+		it("should ensure all messages await initialization promise", async () => {
+			const host = createTestHost()
+			let initializationResolved = false
+			let resolveInit: () => void
+
+			// Create a manually-resolved initialization promise
+			const initPromise = new Promise<void>((resolve) => {
+				resolveInit = () => {
+					initializationResolved = true
+					resolve()
+				}
+			})
+
+			// Override the initialization promise before marking ready
+			const privateHost = host as unknown as { initializationPromise: Promise<void>; isReady: boolean }
+			privateHost.isReady = true
+			privateHost.initializationPromise = initPromise
+
+			const emitSpy = vi.spyOn(host, "emit")
+
+			// Send a message - it should wait for initialization
+			const sendPromise = host.sendToExtension({ type: "requestModes" })
+
+			// Give it a tick to start processing
+			await new Promise((resolve) => setImmediate(resolve))
+
+			// Message should not be emitted yet because initialization hasn't resolved
+			expect(initializationResolved).toBe(false)
+			expect(emitSpy).not.toHaveBeenCalledWith("webviewMessage", { type: "requestModes" })
+
+			// Now resolve initialization
+			resolveInit!()
+
+			// Wait for the message to be sent
+			await sendPromise
+
+			// Now initialization should be complete and message emitted
+			expect(initializationResolved).toBe(true)
+			expect(emitSpy).toHaveBeenCalledWith("webviewMessage", { type: "requestModes" })
+		})
+	})
+
 	describe("public agent state API", () => {
 		it("should return agent state from getAgentState()", () => {
 			const host = createTestHost()
@@ -658,6 +718,8 @@ describe("ExtensionHost", () => {
 			const taskPromise = host.runTask("delegate this").then(() => {
 				settled = true
 			})
+			// Wait for newTask message to be emitted
+			await new Promise((resolve) => setImmediate(resolve))
 			const rootId = (messages.find((message) => message.type === "newTask") as { taskId: string }).taskId
 
 			api.emit(RooCodeEventName.TaskModeSwitched, rootId, "architect")

apps/cli/src/agent/extension-host.ts

@@ -124,7 +124,7 @@ export interface ExtensionHostInterface extends IExtensionHost<ExtensionHostEven
 	runTask(prompt: string, taskId?: string, configuration?: RooCodeSettings, images?: string[]): Promise<void>
 	resumeTask(taskId: string): Promise<void>
 	cancelTask(): Promise<void>
-	sendToExtension(message: WebviewMessage): void
+	sendToExtension(message: WebviewMessage): Promise<void>
 	dispose(): Promise<void>
 }
 
@@ -213,7 +213,9 @@ export class ExtensionHost extends EventEmitter implements ExtensionHostInterfac
 
 		// Initialize client - single source of truth for agent state (including mode).
 		this.client = new ExtensionClient({
-			sendMessage: (msg) => this.sendToExtension(msg),
+			sendMessage: (msg) => {
+				void this.sendToExtension(msg)
+			},
 			debug: options.debug, // Enable debug logging in the client.
 		})
 
@@ -230,7 +232,9 @@ export class ExtensionHost extends EventEmitter implements ExtensionHostInterfac
 		this.askDispatcher = new AskDispatcher({
 			outputManager: this.outputManager,
 			promptManager: this.promptManager,
-			sendMessage: (msg) => this.sendToExtension(msg),
+			sendMessage: (msg) => {
+				void this.sendToExtension(msg)
+			},
 			nonInteractive: options.autonomous || options.nonInteractive,
 			exitOnError: options.exitOnError,
 			disabled: options.disableAskHandling,
@@ -503,11 +507,12 @@ export class ExtensionHost extends EventEmitter implements ExtensionHostInterfac
 	// Message Handling
 	// ==========================================================================
 
-	public sendToExtension(message: WebviewMessage): void {
+	public async sendToExtension(message: WebviewMessage): Promise<void> {
 		if (!this.isReady) {
 			throw new Error("You cannot send messages to the extension before it is ready")
 		}
 
+		await this.initializationPromise
 		this.emit("webviewMessage", message)
 	}
 
@@ -638,7 +643,7 @@ export class ExtensionHost extends EventEmitter implements ExtensionHostInterfac
 		const rootTaskId = taskId ?? (this.options.autonomous ? randomUUID() : "unknown")
 		this.rootTaskId = rootTaskId
 		const completion = this.waitForTaskCompletion(rootTaskId)
-		this.sendToExtension({
+		await this.sendToExtension({
 			type: "newTask",
 			text: prompt,
 			...(taskId || this.options.autonomous ? { taskId: rootTaskId } : {}),
@@ -651,7 +656,7 @@ export class ExtensionHost extends EventEmitter implements ExtensionHostInterfac
 	public async resumeTask(taskId: string): Promise<void> {
 		this.rootTaskId = taskId
 		const completion = this.waitForTaskCompletion(taskId)
-		this.sendToExtension({ type: "showTaskWithId", text: taskId })
+		await this.sendToExtension({ type: "showTaskWithId", text: taskId })
 		this.lastTaskResult = await completion
 	}
 

apps/cli/src/ui/hooks/useExtensionHost.ts

@@ -157,8 +157,8 @@ export function useExtensionHost({
 
 				// Request initial state from extension (triggers
 				// postStateToWebview which includes taskHistory).
-				host.sendToExtension({ type: "requestCommands" })
-				host.sendToExtension({ type: "requestModes" })
+				await host.sendToExtension({ type: "requestCommands" })
+				await host.sendToExtension({ type: "requestModes" })
 
 				if (requestedSessionId || continueSession) {
 					await pWaitFor(() => hasReceivedTaskHistory, {