Some cleanup in ExtensionHost (#10600)

Co-authored-by: Roo Code <[email protected]>

apps/cli/docs/AGENT_LOOP.md

@@ -180,12 +180,13 @@ function isStreaming(messages) {
 
 ### ExtensionClient
 
-The **single source of truth** for agent state. It:
+The **single source of truth** for agent state, including the current mode. It:
 
 - Receives all messages from the extension
 - Stores them in the `StateStore`
+- Tracks the current mode from state messages
 - Computes the current state via `detectAgentState()`
-- Emits events when state changes
+- Emits events when state changes (including mode changes)
 
 ```typescript
 const client = new ExtensionClient({
@@ -199,21 +200,31 @@ if (state.isWaitingForInput) {
 	console.log(`Agent needs: ${state.currentAsk}`)
 }
 
+// Query current mode
+const mode = client.getCurrentMode()
+console.log(`Current mode: ${mode}`) // e.g., "code", "architect", "ask"
+
 // Subscribe to events
 client.on("waitingForInput", (event) => {
 	console.log(`Waiting for: ${event.ask}`)
 })
+
+// Subscribe to mode changes
+client.on("modeChanged", (event) => {
+	console.log(`Mode changed: ${event.previousMode} -> ${event.currentMode}`)
+})
 ```
 
 ### StateStore
 
-Holds the `clineMessages` array and computed state:
+Holds the `clineMessages` array, computed state, and current mode:
 
 ```typescript
 interface StoreState {
 	messages: ClineMessage[] // The raw message array
 	agentState: AgentStateInfo // Computed state
 	isInitialized: boolean // Have we received any state?
+	currentMode: string | undefined // Current mode (e.g., "code", "architect")
 }
 ```
 
@@ -221,9 +232,9 @@ interface StoreState {
 
 Handles incoming messages from the extension:
 
-- `"state"` messages → Update `clineMessages` array
+- `"state"` messages → Update `clineMessages` array and track mode
 - `"messageUpdated"` messages → Update single message in array
-- Emits events for state transitions
+- Emits events for state transitions and mode changes
 
 ### AskDispatcher
 

apps/cli/package.json

@@ -30,6 +30,7 @@
 		"commander": "^12.1.0",
 		"fuzzysort": "^3.1.0",
 		"ink": "^6.6.0",
+		"p-wait-for": "^5.0.2",
 		"react": "^19.1.0",
 		"superjson": "^2.2.6",
 		"zustand": "^5.0.0"

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

@@ -14,8 +14,8 @@ function createMessage(overrides: Partial<ClineMessage>): ClineMessage {
 	return { ts: Date.now() + Math.random() * 1000, type: "say", ...overrides }
 }
 
-function createStateMessage(messages: ClineMessage[]): ExtensionMessage {
-	return { type: "state", state: { clineMessages: messages } } as ExtensionMessage
+function createStateMessage(messages: ClineMessage[], mode?: string): ExtensionMessage {
+	return { type: "state", state: { clineMessages: messages, mode } } as ExtensionMessage
 }
 
 describe("detectAgentState", () => {
@@ -300,6 +300,44 @@ describe("ExtensionClient", () => {
 			client.handleMessage(createStateMessage([createMessage({ say: "text", ts: Date.now() + 1 })]))
 			expect(callCount).toBe(1) // Should not increase.
 		})
+
+		it("should emit modeChanged events", () => {
+			const { client } = createMockClient()
+			const modeChanges: { previousMode: string | undefined; currentMode: string }[] = []
+
+			client.onModeChanged((event) => {
+				modeChanges.push(event)
+			})
+
+			// Set initial mode
+			client.handleMessage(createStateMessage([createMessage({ say: "text" })], "code"))
+
+			expect(modeChanges).toHaveLength(1)
+			expect(modeChanges[0]).toEqual({ previousMode: undefined, currentMode: "code" })
+
+			// Change mode
+			client.handleMessage(createStateMessage([createMessage({ say: "text" })], "architect"))
+
+			expect(modeChanges).toHaveLength(2)
+			expect(modeChanges[1]).toEqual({ previousMode: "code", currentMode: "architect" })
+		})
+
+		it("should not emit modeChanged when mode stays the same", () => {
+			const { client } = createMockClient()
+			let modeChangeCount = 0
+
+			client.onModeChanged(() => {
+				modeChangeCount++
+			})
+
+			// Set initial mode
+			client.handleMessage(createStateMessage([createMessage({ say: "text" })], "code"))
+			expect(modeChangeCount).toBe(1)
+
+			// Same mode - should not emit
+			client.handleMessage(createStateMessage([createMessage({ say: "text", ts: Date.now() + 1 })], "code"))
+			expect(modeChangeCount).toBe(1)
+		})
 	})
 
 	describe("Response methods", () => {
@@ -458,6 +496,65 @@ describe("ExtensionClient", () => {
 			expect(client.isInitialized()).toBe(false)
 			expect(client.getCurrentState()).toBe(AgentLoopState.NO_TASK)
 		})
+
+		it("should reset mode on reset", () => {
+			const { client } = createMockClient()
+
+			client.handleMessage(createStateMessage([createMessage({ say: "text" })], "code"))
+			expect(client.getCurrentMode()).toBe("code")
+
+			client.reset()
+
+			expect(client.getCurrentMode()).toBeUndefined()
+		})
+	})
+
+	describe("Mode tracking", () => {
+		it("should return undefined mode when not initialized", () => {
+			const { client } = createMockClient()
+			expect(client.getCurrentMode()).toBeUndefined()
+		})
+
+		it("should track mode from state messages", () => {
+			const { client } = createMockClient()
+
+			client.handleMessage(createStateMessage([createMessage({ say: "text" })], "code"))
+
+			expect(client.getCurrentMode()).toBe("code")
+		})
+
+		it("should update mode when it changes", () => {
+			const { client } = createMockClient()
+
+			client.handleMessage(createStateMessage([createMessage({ say: "text" })], "code"))
+			expect(client.getCurrentMode()).toBe("code")
+
+			client.handleMessage(createStateMessage([createMessage({ say: "text", ts: Date.now() + 1 })], "architect"))
+			expect(client.getCurrentMode()).toBe("architect")
+		})
+
+		it("should preserve mode when state message has no mode", () => {
+			const { client } = createMockClient()
+
+			// Set initial mode
+			client.handleMessage(createStateMessage([createMessage({ say: "text" })], "code"))
+			expect(client.getCurrentMode()).toBe("code")
+
+			// State update without mode - should preserve existing mode
+			client.handleMessage(createStateMessage([createMessage({ say: "text", ts: Date.now() + 1 })]))
+			expect(client.getCurrentMode()).toBe("code")
+		})
+
+		it("should preserve mode when task is cleared", () => {
+			const { client } = createMockClient()
+
+			client.handleMessage(createStateMessage([createMessage({ say: "text" })], "architect"))
+			expect(client.getCurrentMode()).toBe("architect")
+
+			client.clearTask()
+			// Mode should be preserved after clear
+			expect(client.getCurrentMode()).toBe("architect")
+		})
 	})
 })
 

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

@@ -2,12 +2,12 @@
 
 import { EventEmitter } from "events"
 import fs from "fs"
-import os from "os"
-import path from "path"
 
-import type { WebviewMessage } from "@roo-code/types"
+import type { ExtensionMessage, WebviewMessage } from "@roo-code/types"
 
 import { type ExtensionHostOptions, ExtensionHost } from "../extension-host.js"
+import { ExtensionClient } from "../extension-client.js"
+import { AgentLoopState } from "../agent-state.js"
 
 vi.mock("@roo-code/vscode-shim", () => ({
 	createVSCodeAPI: vi.fn(() => ({
@@ -16,6 +16,10 @@ vi.mock("@roo-code/vscode-shim", () => ({
 	setRuntimeConfigValues: vi.fn(),
 }))
 
+vi.mock("@/lib/storage/index.js", () => ({
+	createEphemeralStorageDir: vi.fn(() => Promise.resolve("/tmp/roo-cli-test-ephemeral")),
+}))
+
 /**
  * Create a test ExtensionHost with default options.
  */
@@ -46,8 +50,16 @@ function getPrivate<T>(host: ExtensionHost, key: string): T {
 	return (host as unknown as PrivateHost)[key] as T
 }
 
+/**
+ * Helper to set private members for testing
+ */
+function setPrivate(host: ExtensionHost, key: string, value: unknown): void {
+	;(host as unknown as PrivateHost)[key] = value
+}
+
 /**
  * Helper to call private methods for testing
+ * This uses a more permissive type to avoid TypeScript errors with private methods
  */
 function callPrivate<T>(host: ExtensionHost, method: string, ...args: unknown[]): T {
 	const fn = (host as unknown as PrivateHost)[method] as ((...a: unknown[]) => T) | undefined
@@ -86,7 +98,12 @@ describe("ExtensionHost", () => {
 
 			const host = new ExtensionHost(options)
 
-			expect(getPrivate(host, "options")).toEqual(options)
+			// Options are stored but integrationTest is set to true
+			const storedOptions = getPrivate<ExtensionHostOptions>(host, "options")
+			expect(storedOptions.mode).toBe(options.mode)
+			expect(storedOptions.workspacePath).toBe(options.workspacePath)
+			expect(storedOptions.extensionPath).toBe(options.extensionPath)
+			expect(storedOptions.integrationTest).toBe(true) // Always set to true in constructor
 		})
 
 		it("should be an EventEmitter instance", () => {
@@ -97,8 +114,7 @@ describe("ExtensionHost", () => {
 		it("should initialize with default state values", () => {
 			const host = createTestHost()
 
-			expect(getPrivate(host, "isWebviewReady")).toBe(false)
-			expect(getPrivate<unknown[]>(host, "pendingMessages")).toEqual([])
+			expect(getPrivate(host, "isReady")).toBe(false)
 			expect(getPrivate(host, "vscode")).toBeNull()
 			expect(getPrivate(host, "extensionModule")).toBeNull()
 		})
@@ -115,25 +131,26 @@ describe("ExtensionHost", () => {
 	})
 
 	describe("webview provider registration", () => {
-		it("should register webview provider", () => {
+		it("should register webview provider without throwing", () => {
 			const host = createTestHost()
 			const mockProvider = { resolveWebviewView: vi.fn() }
 
-			host.registerWebviewProvider("test-view", mockProvider)
-
-			const providers = getPrivate<Map<string, unknown>>(host, "webviewProviders")
-			expect(providers.get("test-view")).toBe(mockProvider)
+			// registerWebviewProvider is now a no-op, just ensure it doesn't throw
+			expect(() => {
+				host.registerWebviewProvider("test-view", mockProvider)
+			}).not.toThrow()
 		})
 
-		it("should unregister webview provider", () => {
+		it("should unregister webview provider without throwing", () => {
 			const host = createTestHost()
 			const mockProvider = { resolveWebviewView: vi.fn() }
 
 			host.registerWebviewProvider("test-view", mockProvider)
-			host.unregisterWebviewProvider("test-view")
 
-			const providers = getPrivate<Map<string, unknown>>(host, "webviewProviders")
-			expect(providers.has("test-view")).toBe(false)
+			// unregisterWebviewProvider is now a no-op, just ensure it doesn't throw
+			expect(() => {
+				host.unregisterWebviewProvider("test-view")
+			}).not.toThrow()
 		})
 
 		it("should handle unregistering non-existent provider gracefully", () => {
@@ -160,49 +177,48 @@ describe("ExtensionHost", () => {
 		})
 
 		describe("markWebviewReady", () => {
-			it("should set isWebviewReady to true", () => {
+			it("should set isReady to true", () => {
 				const host = createTestHost()
 				host.markWebviewReady()
-				expect(getPrivate(host, "isWebviewReady")).toBe(true)
+				expect(getPrivate(host, "isReady")).toBe(true)
 			})
 
-			it("should emit webviewReady event", () => {
+			it("should send webviewDidLaunch message", () => {
 				const host = createTestHost()
-				const listener = vi.fn()
+				const emitSpy = vi.spyOn(host, "emit")
 
-				host.on("webviewReady", listener)
 				host.markWebviewReady()
 
-				expect(listener).toHaveBeenCalled()
+				expect(emitSpy).toHaveBeenCalledWith("webviewMessage", { type: "webviewDidLaunch" })
 			})
 
-			it("should flush pending messages", () => {
+			it("should send updateSettings message", () => {
 				const host = createTestHost()
 				const emitSpy = vi.spyOn(host, "emit")
 
-				// Queue messages before ready
-				host.sendToExtension({ type: "requestModes" })
-				host.sendToExtension({ type: "requestCommands" })
-
-				// Mark ready (should flush)
 				host.markWebviewReady()
 
-				// Check that webviewMessage events were emitted for pending messages
-				expect(emitSpy).toHaveBeenCalledWith("webviewMessage", { type: "requestModes" })
-				expect(emitSpy).toHaveBeenCalledWith("webviewMessage", { type: "requestCommands" })
+				// Check that updateSettings was called
+				const updateSettingsCall = emitSpy.mock.calls.find(
+					(call) =>
+						call[0] === "webviewMessage" &&
+						typeof call[1] === "object" &&
+						call[1] !== null &&
+						(call[1] as WebviewMessage).type === "updateSettings",
+				)
+				expect(updateSettingsCall).toBeDefined()
 			})
 		})
 	})
 
 	describe("sendToExtension", () => {
-		it("should queue message when webview not ready", () => {
+		it("should throw error when extension not ready", () => {
 			const host = createTestHost()
 			const message: WebviewMessage = { type: "requestModes" }
 
-			host.sendToExtension(message)
-
-			const pending = getPrivate<unknown[]>(host, "pendingMessages")
-			expect(pending).toContain(message)
+			expect(() => {
+				host.sendToExtension(message)
+			}).toThrow("You cannot send messages to the extension before it is ready")
 		})
 
 		it("should emit webviewMessage event when webview is ready", () => {
@@ -211,51 +227,37 @@ describe("ExtensionHost", () => {
 			const message: WebviewMessage = { type: "requestModes" }
 
 			host.markWebviewReady()
+			emitSpy.mockClear() // Clear the markWebviewReady calls
 			host.sendToExtension(message)
 
 			expect(emitSpy).toHaveBeenCalledWith("webviewMessage", message)
 		})
 
-		it("should not queue message when webview is ready", () => {
+		it("should not throw when webview is ready", () => {
 			const host = createTestHost()
 
 			host.markWebviewReady()
-			host.sendToExtension({ type: "requestModes" })
 
-			const pending = getPrivate<unknown[]>(host, "pendingMessages")
-			expect(pending).toHaveLength(0)
+			expect(() => {
+				host.sendToExtension({ type: "requestModes" })
+			}).not.toThrow()
 		})
 	})
 
-	describe("handleExtensionMessage", () => {
-		it("should forward messages to the client", () => {
+	describe("message handling via client", () => {
+		it("should forward extension messages to the client", () => {
 			const host = createTestHost()
-			const client = host.getExtensionClient()
-			const handleMessageSpy = vi.spyOn(client, "handleMessage")
+			const client = getPrivate(host, "client") as ExtensionClient
 
-			callPrivate(host, "handleExtensionMessage", { type: "state", state: { clineMessages: [] } })
-
-			expect(handleMessageSpy).toHaveBeenCalled()
-		})
-
-		it("should track mode from state messages", () => {
-			const host = createTestHost()
-
-			callPrivate(host, "handleExtensionMessage", {
+			// Simulate extension message.
+			host.emit("extensionWebviewMessage", {
 				type: "state",
-				state: { mode: "architect", clineMessages: [] },
-			})
+				state: { clineMessages: [] },
+			} as unknown as ExtensionMessage)
 
-			expect(getPrivate(host, "currentMode")).toBe("architect")
-		})
-
-		it("should emit modesUpdated for modes messages", () => {
-			const host = createTestHost()
-			const emitSpy = vi.spyOn(host, "emit")
-
-			callPrivate(host, "handleExtensionMessage", { type: "modes", modes: [] })
-
-			expect(emitSpy).toHaveBeenCalledWith("modesUpdated", { type: "modes", modes: [] })
+			// Message listener is set up in activate(), which we can't easily call in unit tests.
+			// But we can verify the client exists and has the handleMessage method.
+			expect(typeof client.handleMessage).toBe("function")
 		})
 	})
 
@@ -274,94 +276,63 @@ describe("ExtensionHost", () => {
 			const host = createTestHost()
 			expect(typeof host.isWaitingForInput()).toBe("boolean")
 		})
-
-		it("should return isAgentRunning() status", () => {
-			const host = createTestHost()
-			expect(typeof host.isAgentRunning()).toBe("boolean")
-		})
-
-		it("should return the client from getExtensionClient()", () => {
-			const host = createTestHost()
-			const client = host.getExtensionClient()
-
-			expect(client).toBeDefined()
-			expect(typeof client.handleMessage).toBe("function")
-		})
-
-		it("should return the output manager from getOutputManager()", () => {
-			const host = createTestHost()
-			const outputManager = host.getOutputManager()
-
-			expect(outputManager).toBeDefined()
-			expect(typeof outputManager.output).toBe("function")
-		})
-
-		it("should return the prompt manager from getPromptManager()", () => {
-			const host = createTestHost()
-			const promptManager = host.getPromptManager()
-
-			expect(promptManager).toBeDefined()
-		})
-
-		it("should return the ask dispatcher from getAskDispatcher()", () => {
-			const host = createTestHost()
-			const askDispatcher = host.getAskDispatcher()
-
-			expect(askDispatcher).toBeDefined()
-			expect(typeof askDispatcher.handleAsk).toBe("function")
-		})
 	})
 
 	describe("quiet mode", () => {
 		describe("setupQuietMode", () => {
-			it("should suppress console.log, warn, debug, info when enabled", () => {
+			it("should not modify console when integrationTest is true", () => {
+				// By default, constructor sets integrationTest = true
 				const host = createTestHost()
 				const originalLog = console.log
 
 				callPrivate(host, "setupQuietMode")
 
-				// These should be no-ops now (different from original)
-				expect(console.log).not.toBe(originalLog)
-
-				// Verify they are actually no-ops by calling them (should not throw)
-				expect(() => console.log("test")).not.toThrow()
-				expect(() => console.warn("test")).not.toThrow()
-				expect(() => console.debug("test")).not.toThrow()
-				expect(() => console.info("test")).not.toThrow()
-
-				// Restore for other tests
-				callPrivate(host, "restoreConsole")
+				// Console should not be modified since integrationTest is true
+				expect(console.log).toBe(originalLog)
 			})
 
-			it("should preserve console.error", () => {
+			it("should suppress console when integrationTest is false", () => {
 				const host = createTestHost()
-				const originalError = console.error
+				const originalLog = console.log
+
+				// Override integrationTest to false
+				const options = getPrivate<ExtensionHostOptions>(host, "options")
+				options.integrationTest = false
 
 				callPrivate(host, "setupQuietMode")
 
-				expect(console.error).toBe(originalError)
+				// Console should be modified
+				expect(console.log).not.toBe(originalLog)
 
+				// Restore for other tests
 				callPrivate(host, "restoreConsole")
 			})
 
-			it("should store original console methods", () => {
+			it("should preserve console.error even when suppressing", () => {
 				const host = createTestHost()
-				const originalLog = console.log
+				const originalError = console.error
+
+				// Override integrationTest to false
+				const options = getPrivate<ExtensionHostOptions>(host, "options")
+				options.integrationTest = false
 
 				callPrivate(host, "setupQuietMode")
 
-				const stored = getPrivate<{ log: typeof console.log }>(host, "originalConsole")
-				expect(stored.log).toBe(originalLog)
+				expect(console.error).toBe(originalError)
 
 				callPrivate(host, "restoreConsole")
 			})
 		})
 
 		describe("restoreConsole", () => {
-			it("should restore original console methods", () => {
+			it("should restore original console methods when suppressed", () => {
 				const host = createTestHost()
 				const originalLog = console.log
 
+				// Override integrationTest to false to actually suppress
+				const options = getPrivate<ExtensionHostOptions>(host, "options")
+				options.integrationTest = false
+
 				callPrivate(host, "setupQuietMode")
 				callPrivate(host, "restoreConsole")
 
@@ -376,20 +347,6 @@ describe("ExtensionHost", () => {
 				}).not.toThrow()
 			})
 		})
-
-		describe("suppressNodeWarnings", () => {
-			it("should suppress process.emitWarning", () => {
-				const host = createTestHost()
-				const originalEmitWarning = process.emitWarning
-
-				callPrivate(host, "suppressNodeWarnings")
-
-				expect(process.emitWarning).not.toBe(originalEmitWarning)
-
-				// Restore
-				callPrivate(host, "restoreConsole")
-			})
-		})
 	})
 
 	describe("dispose", () => {
@@ -401,7 +358,7 @@ describe("ExtensionHost", () => {
 
 		it("should remove message listener", async () => {
 			const listener = vi.fn()
-			;(host as unknown as Record<string, unknown>).messageListener = listener
+			setPrivate(host, "messageListener", listener)
 			host.on("extensionWebviewMessage", listener)
 
 			await host.dispose()
@@ -411,9 +368,9 @@ describe("ExtensionHost", () => {
 
 		it("should call extension deactivate if available", async () => {
 			const deactivateMock = vi.fn()
-			;(host as unknown as Record<string, unknown>).extensionModule = {
+			setPrivate(host, "extensionModule", {
 				deactivate: deactivateMock,
-			}
+			})
 
 			await host.dispose()
 
@@ -421,7 +378,7 @@ describe("ExtensionHost", () => {
 		})
 
 		it("should clear vscode reference", async () => {
-			;(host as unknown as Record<string, unknown>).vscode = { context: {} }
+			setPrivate(host, "vscode", { context: {} })
 
 			await host.dispose()
 
@@ -429,22 +386,13 @@ describe("ExtensionHost", () => {
 		})
 
 		it("should clear extensionModule reference", async () => {
-			;(host as unknown as Record<string, unknown>).extensionModule = {}
+			setPrivate(host, "extensionModule", {})
 
 			await host.dispose()
 
 			expect(getPrivate(host, "extensionModule")).toBeNull()
 		})
 
-		it("should clear webviewProviders", async () => {
-			host.registerWebviewProvider("test", {})
-
-			await host.dispose()
-
-			const providers = getPrivate<Map<string, unknown>>(host, "webviewProviders")
-			expect(providers.size).toBe(0)
-		})
-
 		it("should delete global vscode", async () => {
 			;(global as Record<string, unknown>).vscode = {}
 
@@ -461,422 +409,188 @@ describe("ExtensionHost", () => {
 			expect((global as Record<string, unknown>).__extensionHost).toBeUndefined()
 		})
 
-		it("should restore console if it was suppressed", async () => {
+		it("should call restoreConsole", async () => {
 			const restoreConsoleSpy = spyOnPrivate(host, "restoreConsole")
 
 			await host.dispose()
 
 			expect(restoreConsoleSpy).toHaveBeenCalled()
 		})
-
-		it("should clear managers", async () => {
-			const outputManager = host.getOutputManager()
-			const askDispatcher = host.getAskDispatcher()
-			const outputClearSpy = vi.spyOn(outputManager, "clear")
-			const askClearSpy = vi.spyOn(askDispatcher, "clear")
-
-			await host.dispose()
-
-			expect(outputClearSpy).toHaveBeenCalled()
-			expect(askClearSpy).toHaveBeenCalled()
-		})
-
-		it("should reset client", async () => {
-			const client = host.getExtensionClient()
-			const resetSpy = vi.spyOn(client, "reset")
-
-			await host.dispose()
-
-			expect(resetSpy).toHaveBeenCalled()
-		})
 	})
 
-	describe("waitForCompletion", () => {
-		it("should resolve when taskComplete is emitted", async () => {
+	describe("runTask", () => {
+		it("should send newTask message when called", async () => {
 			const host = createTestHost()
+			host.markWebviewReady()
 
-			const promise = callPrivate<Promise<void>>(host, "waitForCompletion")
+			const emitSpy = vi.spyOn(host, "emit")
+			const client = getPrivate(host, "client") as ExtensionClient
+
+			// Start the task (will hang waiting for completion)
+			const taskPromise = host.runTask("test prompt")
+
+			// Emit completion to resolve the promise via the client's emitter
+			const taskCompletedEvent = {
+				success: true,
+				stateInfo: {
+					state: AgentLoopState.IDLE,
+					isWaitingForInput: false,
+					isRunning: false,
+					isStreaming: false,
+					requiredAction: "start_task" as const,
+					description: "Task completed",
+				},
+			}
+			setTimeout(() => client.getEmitter().emit("taskCompleted", taskCompletedEvent), 10)
 
-			// Emit completion after a short delay
-			setTimeout(() => host.emit("taskComplete"), 10)
+			await taskPromise
 
-			await expect(promise).resolves.toBeUndefined()
+			expect(emitSpy).toHaveBeenCalledWith("webviewMessage", { type: "newTask", text: "test prompt" })
 		})
 
-		it("should reject when taskError is emitted", async () => {
+		it("should resolve when taskCompleted is emitted on client", async () => {
 			const host = createTestHost()
+			host.markWebviewReady()
 
-			const promise = callPrivate<Promise<void>>(host, "waitForCompletion")
-
-			setTimeout(() => host.emit("taskError", "Test error"), 10)
+			const client = getPrivate(host, "client") as ExtensionClient
+			const taskPromise = host.runTask("test prompt")
+
+			// Emit completion after a short delay via the client's emitter
+			const taskCompletedEvent = {
+				success: true,
+				stateInfo: {
+					state: AgentLoopState.IDLE,
+					isWaitingForInput: false,
+					isRunning: false,
+					isStreaming: false,
+					requiredAction: "start_task" as const,
+					description: "Task completed",
+				},
+			}
+			setTimeout(() => client.getEmitter().emit("taskCompleted", taskCompletedEvent), 10)
 
-			await expect(promise).rejects.toThrow("Test error")
+			await expect(taskPromise).resolves.toBeUndefined()
 		})
 	})
 
-	describe("mode tracking via handleExtensionMessage", () => {
-		let host: ExtensionHost
-
-		beforeEach(() => {
-			host = createTestHost({
-				mode: "code",
-				provider: "anthropic",
-				apiKey: "test-key",
-				model: "test-model",
-			})
-			// Mock process.stdout.write which is used by output()
-			vi.spyOn(process.stdout, "write").mockImplementation(() => true)
-		})
+	describe("initial settings", () => {
+		it("should set mode from options", () => {
+			const host = createTestHost({ mode: "architect" })
 
-		afterEach(() => {
-			vi.restoreAllMocks()
+			const initialSettings = getPrivate<Record<string, unknown>>(host, "initialSettings")
+			expect(initialSettings.mode).toBe("architect")
 		})
 
-		it("should track current mode when state updates with a mode", () => {
-			// Initial state update establishes current mode
-			callPrivate(host, "handleExtensionMessage", { type: "state", state: { mode: "code", clineMessages: [] } })
-			expect(getPrivate(host, "currentMode")).toBe("code")
+		it("should enable auto-approval in non-interactive mode", () => {
+			const host = createTestHost({ nonInteractive: true })
 
-			// Second state update should update tracked mode
-			callPrivate(host, "handleExtensionMessage", {
-				type: "state",
-				state: { mode: "architect", clineMessages: [] },
-			})
-			expect(getPrivate(host, "currentMode")).toBe("architect")
+			const initialSettings = getPrivate<Record<string, unknown>>(host, "initialSettings")
+			expect(initialSettings.autoApprovalEnabled).toBe(true)
+			expect(initialSettings.alwaysAllowReadOnly).toBe(true)
+			expect(initialSettings.alwaysAllowWrite).toBe(true)
+			expect(initialSettings.alwaysAllowExecute).toBe(true)
 		})
 
-		it("should not change current mode when state has no mode", () => {
-			// Initial state update establishes current mode
-			callPrivate(host, "handleExtensionMessage", { type: "state", state: { mode: "code", clineMessages: [] } })
-			expect(getPrivate(host, "currentMode")).toBe("code")
+		it("should disable auto-approval in interactive mode", () => {
+			const host = createTestHost({ nonInteractive: false })
 
-			// State without mode should not change tracked mode
-			callPrivate(host, "handleExtensionMessage", { type: "state", state: { clineMessages: [] } })
-			expect(getPrivate(host, "currentMode")).toBe("code")
+			const initialSettings = getPrivate<Record<string, unknown>>(host, "initialSettings")
+			expect(initialSettings.autoApprovalEnabled).toBe(false)
 		})
 
-		it("should track current mode across multiple changes", () => {
-			// Start with code mode
-			callPrivate(host, "handleExtensionMessage", { type: "state", state: { mode: "code", clineMessages: [] } })
-			expect(getPrivate(host, "currentMode")).toBe("code")
-
-			// Change to architect
-			callPrivate(host, "handleExtensionMessage", {
-				type: "state",
-				state: { mode: "architect", clineMessages: [] },
-			})
-			expect(getPrivate(host, "currentMode")).toBe("architect")
-
-			// Change to debug
-			callPrivate(host, "handleExtensionMessage", { type: "state", state: { mode: "debug", clineMessages: [] } })
-			expect(getPrivate(host, "currentMode")).toBe("debug")
-
-			// Another state update with debug
-			callPrivate(host, "handleExtensionMessage", { type: "state", state: { mode: "debug", clineMessages: [] } })
-			expect(getPrivate(host, "currentMode")).toBe("debug")
-		})
-
-		it("should not send updateSettings on mode change (CLI settings are applied once during runTask)", () => {
-			// This test ensures mode changes don't trigger automatic re-application of API settings.
-			// CLI settings are applied once during runTask() via updateSettings.
-			// Mode-specific provider profiles are handled by the extension's handleModeSwitch.
-			const sendToExtensionSpy = vi.spyOn(host, "sendToExtension")
-
-			// Initial state
-			callPrivate(host, "handleExtensionMessage", { type: "state", state: { mode: "code", clineMessages: [] } })
-			sendToExtensionSpy.mockClear()
-
-			// Mode change should NOT trigger sendToExtension
-			callPrivate(host, "handleExtensionMessage", {
-				type: "state",
-				state: { mode: "architect", clineMessages: [] },
-			})
-			expect(sendToExtensionSpy).not.toHaveBeenCalled()
-		})
-	})
-
-	describe("applyRuntimeSettings - mode switching", () => {
-		it("should use currentMode when set (from user mode switches)", () => {
-			const host = createTestHost({
-				mode: "code", // Initial mode from CLI options
-				provider: "anthropic",
-				apiKey: "test-key",
-				model: "test-model",
-			})
-
-			// Simulate user switching mode via Ctrl+M - this updates currentMode
-			;(host as unknown as Record<string, unknown>).currentMode = "architect"
-
-			// Create settings object to be modified
-			const settings: Record<string, unknown> = {}
-			callPrivate(host, "applyRuntimeSettings", settings)
+		it("should set reasoning effort when specified", () => {
+			const host = createTestHost({ reasoningEffort: "high" })
 
-			// Should use currentMode (architect), not options.mode (code)
-			expect(settings.mode).toBe("architect")
+			const initialSettings = getPrivate<Record<string, unknown>>(host, "initialSettings")
+			expect(initialSettings.enableReasoningEffort).toBe(true)
+			expect(initialSettings.reasoningEffort).toBe("high")
 		})
 
-		it("should fall back to options.mode when currentMode is not set", () => {
-			const host = createTestHost({
-				mode: "code",
-				provider: "anthropic",
-				apiKey: "test-key",
-				model: "test-model",
-			})
+		it("should disable reasoning effort when set to disabled", () => {
+			const host = createTestHost({ reasoningEffort: "disabled" })
 
-			// currentMode is not set (still null from constructor)
-			expect(getPrivate(host, "currentMode")).toBe("code") // Set from options.mode in constructor
-
-			const settings: Record<string, unknown> = {}
-			callPrivate(host, "applyRuntimeSettings", settings)
-
-			// Should use options.mode as fallback
-			expect(settings.mode).toBe("code")
+			const initialSettings = getPrivate<Record<string, unknown>>(host, "initialSettings")
+			expect(initialSettings.enableReasoningEffort).toBe(false)
 		})
 
-		it("should use currentMode even when it differs from initial options.mode", () => {
-			const host = createTestHost({
-				mode: "code",
-				provider: "anthropic",
-				apiKey: "test-key",
-				model: "test-model",
-			})
-
-			// Simulate multiple mode switches: code -> architect -> debug
-			;(host as unknown as Record<string, unknown>).currentMode = "debug"
+		it("should not set reasoning effort when unspecified", () => {
+			const host = createTestHost({ reasoningEffort: "unspecified" })
 
-			const settings: Record<string, unknown> = {}
-			callPrivate(host, "applyRuntimeSettings", settings)
-
-			// Should use the latest currentMode
-			expect(settings.mode).toBe("debug")
+			const initialSettings = getPrivate<Record<string, unknown>>(host, "initialSettings")
+			expect(initialSettings.enableReasoningEffort).toBeUndefined()
+			expect(initialSettings.reasoningEffort).toBeUndefined()
 		})
+	})
 
-		it("should not set mode if neither currentMode nor options.mode is set", () => {
-			const host = createTestHost({
-				// No mode specified - mode defaults to "code" in createTestHost
-				provider: "anthropic",
-				apiKey: "test-key",
-				model: "test-model",
-			})
+	describe("ephemeral mode", () => {
+		it("should store ephemeral option correctly", () => {
+			const host = createTestHost({ ephemeral: true })
 
-			// Explicitly set currentMode to null (edge case)
-			;(host as unknown as Record<string, unknown>).currentMode = null
-			// Also clear options.mode
 			const options = getPrivate<ExtensionHostOptions>(host, "options")
-			options.mode = ""
-
-			const settings: Record<string, unknown> = {}
-			callPrivate(host, "applyRuntimeSettings", settings)
-
-			// Mode should not be set
-			expect(settings.mode).toBeUndefined()
+			expect(options.ephemeral).toBe(true)
 		})
-	})
 
-	describe("mode switching - end to end simulation", () => {
-		let host: ExtensionHost
+		it("should default ephemeralStorageDir to null", () => {
+			const host = createTestHost()
 
-		beforeEach(() => {
-			host = createTestHost({
-				mode: "code",
-				provider: "anthropic",
-				apiKey: "test-key",
-				model: "test-model",
-			})
-			vi.spyOn(process.stdout, "write").mockImplementation(() => true)
+			expect(getPrivate(host, "ephemeralStorageDir")).toBeNull()
 		})
 
-		afterEach(() => {
-			vi.restoreAllMocks()
-		})
+		it("should clean up ephemeral storage directory on dispose", async () => {
+			const host = createTestHost({ ephemeral: true })
 
-		it("should preserve mode switch when starting a new task", () => {
-			// Step 1: Initial state from extension (like webviewDidLaunch response)
-			callPrivate(host, "handleExtensionMessage", {
-				type: "state",
-				state: { mode: "code", clineMessages: [] },
-			})
-			expect(getPrivate(host, "currentMode")).toBe("code")
+			// Set up a mock ephemeral storage directory
+			const mockEphemeralDir = "/tmp/roo-cli-test-ephemeral-cleanup"
+			setPrivate(host, "ephemeralStorageDir", mockEphemeralDir)
 
-			// Step 2: User presses Ctrl+M to switch mode, extension sends new state
-			callPrivate(host, "handleExtensionMessage", {
-				type: "state",
-				state: { mode: "architect", clineMessages: [] },
-			})
-			expect(getPrivate(host, "currentMode")).toBe("architect")
+			// Mock fs.promises.rm
+			const rmMock = vi.spyOn(fs.promises, "rm").mockResolvedValue(undefined)
 
-			// Step 3: When runTask is called, applyRuntimeSettings should use architect
-			const settings: Record<string, unknown> = {}
-			callPrivate(host, "applyRuntimeSettings", settings)
-			expect(settings.mode).toBe("architect")
-		})
+			await host.dispose()
 
-		it("should handle mode switch before any state messages", () => {
-			// currentMode is initialized to options.mode in constructor
-			expect(getPrivate(host, "currentMode")).toBe("code")
+			expect(rmMock).toHaveBeenCalledWith(mockEphemeralDir, { recursive: true, force: true })
+			expect(getPrivate(host, "ephemeralStorageDir")).toBeNull()
 
-			// Without any state messages, should still use options.mode
-			const settings: Record<string, unknown> = {}
-			callPrivate(host, "applyRuntimeSettings", settings)
-			expect(settings.mode).toBe("code")
+			rmMock.mockRestore()
 		})
 
-		it("should track multiple mode switches correctly", () => {
-			// Switch through multiple modes
-			callPrivate(host, "handleExtensionMessage", {
-				type: "state",
-				state: { mode: "code", clineMessages: [] },
-			})
-			callPrivate(host, "handleExtensionMessage", {
-				type: "state",
-				state: { mode: "architect", clineMessages: [] },
-			})
-			callPrivate(host, "handleExtensionMessage", {
-				type: "state",
-				state: { mode: "debug", clineMessages: [] },
-			})
-			callPrivate(host, "handleExtensionMessage", {
-				type: "state",
-				state: { mode: "ask", clineMessages: [] },
-			})
+		it("should not clean up when ephemeralStorageDir is null", async () => {
+			const host = createTestHost()
 
-			// Should use the most recent mode
-			expect(getPrivate(host, "currentMode")).toBe("ask")
+			// ephemeralStorageDir is null by default
+			expect(getPrivate(host, "ephemeralStorageDir")).toBeNull()
 
-			const settings: Record<string, unknown> = {}
-			callPrivate(host, "applyRuntimeSettings", settings)
-			expect(settings.mode).toBe("ask")
-		})
-	})
+			const rmMock = vi.spyOn(fs.promises, "rm").mockResolvedValue(undefined)
 
-	describe("ephemeral mode", () => {
-		describe("constructor", () => {
-			it("should store ephemeral option", () => {
-				const host = createTestHost({ ephemeral: true })
-				const options = getPrivate<ExtensionHostOptions>(host, "options")
-				expect(options.ephemeral).toBe(true)
-			})
+			await host.dispose()
 
-			it("should default ephemeral to undefined", () => {
-				const host = createTestHost()
-				const options = getPrivate<ExtensionHostOptions>(host, "options")
-				expect(options.ephemeral).toBeUndefined()
-			})
+			// rm should not be called when there's no ephemeral storage
+			expect(rmMock).not.toHaveBeenCalled()
 
-			it("should initialize ephemeralStorageDir to null", () => {
-				const host = createTestHost({ ephemeral: true })
-				expect(getPrivate(host, "ephemeralStorageDir")).toBeNull()
-			})
+			rmMock.mockRestore()
 		})
 
-		describe("createEphemeralStorageDir", () => {
-			let createdDirs: string[] = []
+		it("should handle ephemeral storage cleanup errors gracefully", async () => {
+			const host = createTestHost({ ephemeral: true })
 
-			afterEach(async () => {
-				// Clean up any directories created during tests
-				for (const dir of createdDirs) {
-					try {
-						await fs.promises.rm(dir, { recursive: true, force: true })
-					} catch {
-						// Ignore cleanup errors
-					}
-				}
-				createdDirs = []
-			})
+			// Set up a mock ephemeral storage directory
+			setPrivate(host, "ephemeralStorageDir", "/tmp/roo-cli-test-ephemeral-error")
 
-			it("should create a directory in the system temp folder", async () => {
-				const host = createTestHost({ ephemeral: true })
-				const tmpDir = await callPrivate<Promise<string>>(host, "createEphemeralStorageDir")
-				createdDirs.push(tmpDir)
+			// Mock fs.promises.rm to throw an error
+			const rmMock = vi.spyOn(fs.promises, "rm").mockRejectedValue(new Error("Cleanup failed"))
 
-				expect(tmpDir).toContain(os.tmpdir())
-				expect(tmpDir).toContain("roo-cli-")
-				expect(fs.existsSync(tmpDir)).toBe(true)
-			})
-
-			it("should create a unique directory each time", async () => {
-				const host = createTestHost({ ephemeral: true })
-				const dir1 = await callPrivate<Promise<string>>(host, "createEphemeralStorageDir")
-				const dir2 = await callPrivate<Promise<string>>(host, "createEphemeralStorageDir")
-				createdDirs.push(dir1, dir2)
-
-				expect(dir1).not.toBe(dir2)
-				expect(fs.existsSync(dir1)).toBe(true)
-				expect(fs.existsSync(dir2)).toBe(true)
-			})
+			// dispose should not throw even if cleanup fails
+			await expect(host.dispose()).resolves.toBeUndefined()
 
-			it("should include timestamp and random id in directory name", async () => {
-				const host = createTestHost({ ephemeral: true })
-				const tmpDir = await callPrivate<Promise<string>>(host, "createEphemeralStorageDir")
-				createdDirs.push(tmpDir)
-
-				const dirName = path.basename(tmpDir)
-				// Format: roo-cli-{timestamp}-{randomId}
-				expect(dirName).toMatch(/^roo-cli-\d+-[a-z0-9]+$/)
-			})
+			rmMock.mockRestore()
 		})
 
-		describe("dispose - ephemeral cleanup", () => {
-			it("should clean up ephemeral storage directory on dispose", async () => {
-				const host = createTestHost({ ephemeral: true })
-
-				// Create the ephemeral directory
-				const tmpDir = await callPrivate<Promise<string>>(host, "createEphemeralStorageDir")
-				;(host as unknown as Record<string, unknown>).ephemeralStorageDir = tmpDir
-
-				// Verify directory exists
-				expect(fs.existsSync(tmpDir)).toBe(true)
-
-				// Dispose the host
-				await host.dispose()
-
-				// Directory should be removed
-				expect(fs.existsSync(tmpDir)).toBe(false)
-				expect(getPrivate(host, "ephemeralStorageDir")).toBeNull()
-			})
-
-			it("should not fail dispose if ephemeral directory doesn't exist", async () => {
-				const host = createTestHost({ ephemeral: true })
-
-				// Set a non-existent directory
-				;(host as unknown as Record<string, unknown>).ephemeralStorageDir = "/non/existent/path/roo-cli-test"
+		it("should not affect normal mode when ephemeral is false", () => {
+			const host = createTestHost({ ephemeral: false })
 
-				// Dispose should not throw
-				await expect(host.dispose()).resolves.toBeUndefined()
-			})
-
-			it("should clean up ephemeral directory with contents", async () => {
-				const host = createTestHost({ ephemeral: true })
-
-				// Create the ephemeral directory with some content
-				const tmpDir = await callPrivate<Promise<string>>(host, "createEphemeralStorageDir")
-				;(host as unknown as Record<string, unknown>).ephemeralStorageDir = tmpDir
-
-				// Add some files and subdirectories
-				await fs.promises.writeFile(path.join(tmpDir, "test.txt"), "test content")
-				await fs.promises.mkdir(path.join(tmpDir, "subdir"))
-				await fs.promises.writeFile(path.join(tmpDir, "subdir", "nested.txt"), "nested content")
-
-				// Verify content exists
-				expect(fs.existsSync(path.join(tmpDir, "test.txt"))).toBe(true)
-				expect(fs.existsSync(path.join(tmpDir, "subdir", "nested.txt"))).toBe(true)
-
-				// Dispose the host
-				await host.dispose()
-
-				// Directory and all contents should be removed
-				expect(fs.existsSync(tmpDir)).toBe(false)
-			})
-
-			it("should not clean up anything if not in ephemeral mode", async () => {
-				const host = createTestHost({ ephemeral: false })
-
-				// ephemeralStorageDir should be null
-				expect(getPrivate(host, "ephemeralStorageDir")).toBeNull()
-
-				// Dispose should complete normally
-				await expect(host.dispose()).resolves.toBeUndefined()
-			})
+			const options = getPrivate<ExtensionHostOptions>(host, "options")
+			expect(options.ephemeral).toBe(false)
+			expect(getPrivate(host, "ephemeralStorageDir")).toBeNull()
 		})
 	})
 })

apps/cli/src/agent/events.ts

@@ -71,6 +71,11 @@ export interface ClientEventMap {
 	 */
 	taskCleared: void
 
+	/**
+	 * Emitted when the current mode changes.
+	 */
+	modeChanged: ModeChangedEvent
+
 	/**
 	 * Emitted on any error during message processing.
 	 */
@@ -113,6 +118,16 @@ export interface TaskCompletedEvent {
 	message?: ClineMessage
 }
 
+/**
+ * Event payload when mode changes.
+ */
+export interface ModeChangedEvent {
+	/** The previous mode (undefined if first mode set) */
+	previousMode: string | undefined
+	/** The new/current mode */
+	currentMode: string
+}
+
 // =============================================================================
 // Typed Event Emitter
 // =============================================================================

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

@@ -36,6 +36,7 @@ import {
 	type ClientEventMap,
 	type AgentStateChangeEvent,
 	type WaitingForInputEvent,
+	type ModeChangedEvent,
 } from "./events.js"
 import { AgentLoopState, type AgentStateInfo } from "./agent-state.js"
 
@@ -154,10 +155,12 @@ export class ExtensionClient {
 
 		if (typeof message === "string") {
 			parsed = parseExtensionMessage(message)
+
 			if (!parsed) {
 				if (this.debug) {
 					console.log("[ExtensionClient] Failed to parse message:", message)
 				}
+
 				return
 			}
 		} else {
@@ -257,6 +260,14 @@ export class ExtensionClient {
 		return this.store.isInitialized()
 	}
 
+	/**
+	 * Get the current mode (e.g., "code", "architect", "ask").
+	 * Returns undefined if no mode has been received yet.
+	 */
+	getCurrentMode(): string | undefined {
+		return this.store.getCurrentMode()
+	}
+
 	// ===========================================================================
 	// Event Subscriptions - Realtime notifications
 	// ===========================================================================
@@ -319,6 +330,13 @@ export class ExtensionClient {
 		return this.on("waitingForInput", listener)
 	}
 
+	/**
+	 * Convenience method: Subscribe only to mode changes.
+	 */
+	onModeChanged(listener: (event: ModeChangedEvent) => void): () => void {
+		return this.on("modeChanged", listener)
+	}
+
 	// ===========================================================================
 	// Response Methods - Send actions to the extension
 	// ===========================================================================

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

@@ -6,20 +6,15 @@
  * 2. Loading the extension bundle via require()
  * 3. Activating the extension
  * 4. Wiring up managers for output, prompting, and ask handling
- *
- * Managers handle all the heavy lifting:
- * - ExtensionClient: Agent state detection (single source of truth)
- * - OutputManager: CLI output and streaming
- * - PromptManager: User input collection
- * - AskDispatcher: Ask routing and handling
  */
 
-import { EventEmitter } from "events"
 import { createRequire } from "module"
 import path from "path"
 import { fileURLToPath } from "url"
 import fs from "fs"
-import os from "os"
+import { EventEmitter } from "events"
+
+import pWaitFor from "p-wait-for"
 
 import type {
 	ClineMessage,
@@ -28,15 +23,16 @@ import type {
 	RooCodeSettings,
 	WebviewMessage,
 } from "@roo-code/types"
-import { createVSCodeAPI, setRuntimeConfigValues } from "@roo-code/vscode-shim"
+import { createVSCodeAPI, IExtensionHost, ExtensionHostEventMap, setRuntimeConfigValues } from "@roo-code/vscode-shim"
 import { DebugLogger } from "@roo-code/core/cli"
 
 import type { SupportedProvider } from "@/types/index.js"
 import type { User } from "@/lib/sdk/index.js"
 import { getProviderSettings } from "@/lib/utils/provider.js"
+import { createEphemeralStorageDir } from "@/lib/storage/index.js"
 
-import type { AgentStateChangeEvent, WaitingForInputEvent, TaskCompletedEvent } from "./events.js"
-import { type AgentStateInfo, AgentLoopState } from "./agent-state.js"
+import type { WaitingForInputEvent, TaskCompletedEvent } from "./events.js"
+import type { AgentStateInfo } from "./agent-state.js"
 import { ExtensionClient } from "./extension-client.js"
 import { OutputManager } from "./output-manager.js"
 import { PromptManager } from "./prompt-manager.js"
@@ -52,10 +48,6 @@ const cliLogger = new DebugLogger("CLI")
 const __dirname = path.dirname(fileURLToPath(import.meta.url))
 const CLI_PACKAGE_ROOT = process.env.ROO_CLI_ROOT || path.resolve(__dirname, "..")
 
-// =============================================================================
-// Types
-// =============================================================================
-
 export interface ExtensionHostOptions {
 	mode: string
 	reasoningEffort?: ReasoningEffortExtended | "unspecified" | "disabled"
@@ -92,22 +84,25 @@ interface WebviewViewProvider {
 	resolveWebviewView?(webviewView: unknown, context: unknown, token: unknown): void | Promise<void>
 }
 
-// =============================================================================
-// ExtensionHost Class
-// =============================================================================
+export interface ExtensionHostInterface extends IExtensionHost<ExtensionHostEventMap> {
+	client: ExtensionClient
+	activate(): Promise<void>
+	runTask(prompt: string): Promise<void>
+	sendToExtension(message: WebviewMessage): void
+	dispose(): Promise<void>
+}
 
-export class ExtensionHost extends EventEmitter {
-	// Extension lifecycle
+export class ExtensionHost extends EventEmitter implements ExtensionHostInterface {
+	// Extension lifecycle.
 	private vscode: ReturnType<typeof createVSCodeAPI> | null = null
 	private extensionModule: ExtensionModule | null = null
 	private extensionAPI: unknown = null
-	private webviewProviders: Map<string, WebviewViewProvider> = new Map()
 	private options: ExtensionHostOptions
-	private isWebviewReady = false
-	private pendingMessages: unknown[] = []
+	private isReady = false
 	private messageListener: ((message: ExtensionMessage) => void) | null = null
+	private initialSettings: RooCodeSettings
 
-	// Console suppression
+	// Console suppression.
 	private originalConsole: {
 		log: typeof console.log
 		warn: typeof console.warn
@@ -115,12 +110,10 @@ export class ExtensionHost extends EventEmitter {
 		debug: typeof console.debug
 		info: typeof console.info
 	} | null = null
-	private originalProcessEmitWarning: typeof process.emitWarning | null = null
 
-	// Mode tracking
-	private currentMode: string | null = null
+	private originalProcessEmitWarning: typeof process.emitWarning | null = null
 
-	// Ephemeral storage
+	// Ephemeral storage.
 	private ephemeralStorageDir: string | null = null
 
 	// ==========================================================================
@@ -131,7 +124,7 @@ export class ExtensionHost extends EventEmitter {
 	 * ExtensionClient: Single source of truth for agent loop state.
 	 * Handles message processing and state detection.
 	 */
-	private client: ExtensionClient
+	public readonly client: ExtensionClient
 
 	/**
 	 * OutputManager: Handles all CLI output and streaming.
@@ -159,9 +152,9 @@ export class ExtensionHost extends EventEmitter {
 		super()
 
 		this.options = options
-		this.currentMode = options.mode || null
+		this.options.integrationTest = true
 
-		// Initialize client - single source of truth for agent state.
+		// Initialize client - single source of truth for agent state (including mode).
 		this.client = new ExtensionClient({
 			sendMessage: (msg) => this.sendToExtension(msg),
 			debug: options.debug, // Enable debug logging in the client.
@@ -189,6 +182,47 @@ export class ExtensionHost extends EventEmitter {
 
 		// Wire up client events.
 		this.setupClientEventHandlers()
+
+		// Populate initial settings.
+		const baseSettings: RooCodeSettings = {
+			mode: this.options.mode,
+			commandExecutionTimeout: 30,
+			browserToolEnabled: false,
+			enableCheckpoints: false,
+			...getProviderSettings(this.options.provider, this.options.apiKey, this.options.model),
+		}
+
+		this.initialSettings = this.options.nonInteractive
+			? {
+					autoApprovalEnabled: true,
+					alwaysAllowReadOnly: true,
+					alwaysAllowReadOnlyOutsideWorkspace: true,
+					alwaysAllowWrite: true,
+					alwaysAllowWriteOutsideWorkspace: true,
+					alwaysAllowWriteProtected: true,
+					alwaysAllowBrowser: true,
+					alwaysAllowMcp: true,
+					alwaysAllowModeSwitch: true,
+					alwaysAllowSubtasks: true,
+					alwaysAllowExecute: true,
+					allowedCommands: ["*"],
+					...baseSettings,
+				}
+			: {
+					autoApprovalEnabled: false,
+					...baseSettings,
+				}
+
+		if (this.options.reasoningEffort && this.options.reasoningEffort !== "unspecified") {
+			if (this.options.reasoningEffort === "disabled") {
+				this.initialSettings.enableReasoningEffort = false
+			} else {
+				this.initialSettings.enableReasoningEffort = true
+				this.initialSettings.reasoningEffort = this.options.reasoningEffort
+			}
+		}
+
+		this.setupQuietMode()
 	}
 
 	// ==========================================================================
@@ -200,11 +234,6 @@ export class ExtensionHost extends EventEmitter {
 	 * The client emits events, managers handle them.
 	 */
 	private setupClientEventHandlers(): void {
-		// Forward state changes for external consumers.
-		this.client.on("stateChange", (event: AgentStateChangeEvent) => {
-			this.emit("agentStateChange", event)
-		})
-
 		// Handle new messages - delegate to OutputManager.
 		this.client.on("message", (msg: ClineMessage) => {
 			this.logMessageDebug(msg, "new")
@@ -219,61 +248,34 @@ export class ExtensionHost extends EventEmitter {
 
 		// Handle waiting for input - delegate to AskDispatcher.
 		this.client.on("waitingForInput", (event: WaitingForInputEvent) => {
-			this.emit("agentWaitingForInput", event)
 			this.askDispatcher.handleAsk(event.message)
 		})
 
 		// Handle task completion.
 		this.client.on("taskCompleted", (event: TaskCompletedEvent) => {
-			this.emit("agentTaskCompleted", event)
-			this.handleTaskCompleted(event)
-		})
-	}
-
-	/**
-	 * Debug logging for messages (first/last pattern).
-	 */
-	private logMessageDebug(msg: ClineMessage, type: "new" | "updated"): void {
-		if (msg.partial) {
-			if (!this.outputManager.hasLoggedFirstPartial(msg.ts)) {
-				this.outputManager.setLoggedFirstPartial(msg.ts)
-				cliLogger.debug("message:start", { ts: msg.ts, type: msg.say || msg.ask })
+			// Output completion message via OutputManager.
+			// Note: completion_result is an "ask" type, not a "say" type.
+			if (event.message && event.message.type === "ask" && event.message.ask === "completion_result") {
+				this.outputManager.outputCompletionResult(event.message.ts, event.message.text || "")
 			}
-		} else {
-			cliLogger.debug(`message:${type === "new" ? "new" : "complete"}`, { ts: msg.ts, type: msg.say || msg.ask })
-			this.outputManager.clearLoggedFirstPartial(msg.ts)
-		}
-	}
-
-	/**
-	 * Handle task completion.
-	 */
-	private handleTaskCompleted(event: TaskCompletedEvent): void {
-		// Output completion message via OutputManager.
-		// Note: completion_result is an "ask" type, not a "say" type.
-		if (event.message && event.message.type === "ask" && event.message.ask === "completion_result") {
-			this.outputManager.outputCompletionResult(event.message.ts, event.message.text || "")
-		}
-
-		// Emit taskComplete for waitForCompletion.
-		this.emit("taskComplete")
+		})
 	}
 
 	// ==========================================================================
-	// Console Suppression
+	// Logging + Console Suppression
 	// ==========================================================================
 
-	private suppressNodeWarnings(): void {
-		this.originalProcessEmitWarning = process.emitWarning
-		process.emitWarning = () => {}
-		process.on("warning", () => {})
-	}
-
 	private setupQuietMode(): void {
 		if (this.options.integrationTest) {
 			return
 		}
 
+		// Suppress node warnings.
+		this.originalProcessEmitWarning = process.emitWarning
+		process.emitWarning = () => {}
+		process.on("warning", () => {})
+
+		// Suppress console output.
 		this.originalConsole = {
 			log: console.log,
 			warn: console.warn,
@@ -308,21 +310,23 @@ export class ExtensionHost extends EventEmitter {
 		}
 	}
 
+	private logMessageDebug(msg: ClineMessage, type: "new" | "updated"): void {
+		if (msg.partial) {
+			if (!this.outputManager.hasLoggedFirstPartial(msg.ts)) {
+				this.outputManager.setLoggedFirstPartial(msg.ts)
+				cliLogger.debug("message:start", { ts: msg.ts, type: msg.say || msg.ask })
+			}
+		} else {
+			cliLogger.debug(`message:${type === "new" ? "new" : "complete"}`, { ts: msg.ts, type: msg.say || msg.ask })
+			this.outputManager.clearLoggedFirstPartial(msg.ts)
+		}
+	}
+
 	// ==========================================================================
 	// Extension Lifecycle
 	// ==========================================================================
 
-	private async createEphemeralStorageDir(): Promise<string> {
-		const uniqueId = `${Date.now()}-${Math.random().toString(36).slice(2, 11)}`
-		const tmpDir = path.join(os.tmpdir(), `roo-cli-${uniqueId}`)
-		await fs.promises.mkdir(tmpDir, { recursive: true })
-		return tmpDir
-	}
-
-	async activate(): Promise<void> {
-		this.suppressNodeWarnings()
-		this.setupQuietMode()
-
+	public async activate(): Promise<void> {
 		const bundlePath = path.join(this.options.extensionPath, "extension.js")
 
 		if (!fs.existsSync(bundlePath)) {
@@ -333,8 +337,8 @@ export class ExtensionHost extends EventEmitter {
 		let storageDir: string | undefined
 
 		if (this.options.ephemeral) {
-			storageDir = await this.createEphemeralStorageDir()
-			this.ephemeralStorageDir = storageDir
+			this.ephemeralStorageDir = await createEphemeralStorageDir()
+			storageDir = this.ephemeralStorageDir
 		}
 
 		// Create VSCode API mock.
@@ -372,6 +376,7 @@ export class ExtensionHost extends EventEmitter {
 			this.extensionModule = require(bundlePath) as ExtensionModule
 		} catch (error) {
 			Module._resolveFilename = originalResolve
+
 			throw new Error(
 				`Failed to load extension bundle: ${error instanceof Error ? error.message : String(error)}`,
 			)
@@ -385,234 +390,106 @@ export class ExtensionHost extends EventEmitter {
 			throw new Error(`Failed to activate extension: ${error instanceof Error ? error.message : String(error)}`)
 		}
 
-		// Set up message listener - forward all messages to client
-		this.messageListener = (message: ExtensionMessage) => this.handleExtensionMessage(message)
+		// Set up message listener - forward all messages to client.
+		this.messageListener = (message: ExtensionMessage) => this.client.handleMessage(message)
 		this.on("extensionWebviewMessage", this.messageListener)
+
+		await pWaitFor(() => this.isReady, { interval: 100, timeout: 10_000 })
 	}
 
-	// ==========================================================================
-	// Webview Provider Registration
-	// ==========================================================================
+	public registerWebviewProvider(_viewId: string, _provider: WebviewViewProvider): void {}
 
-	registerWebviewProvider(viewId: string, provider: WebviewViewProvider): void {
-		this.webviewProviders.set(viewId, provider)
-	}
+	public unregisterWebviewProvider(_viewId: string): void {}
 
-	unregisterWebviewProvider(viewId: string): void {
-		this.webviewProviders.delete(viewId)
-	}
+	public markWebviewReady(): void {
+		this.isReady = true
 
-	isInInitialSetup(): boolean {
-		return !this.isWebviewReady
-	}
+		// Send initial webview messages to trigger proper extension initialization.
+		// This is critical for the extension to start sending state updates properly.
+		this.sendToExtension({ type: "webviewDidLaunch" })
 
-	markWebviewReady(): void {
-		this.isWebviewReady = true
-		this.emit("webviewReady")
-		this.flushPendingMessages()
+		setRuntimeConfigValues("roo-cline", this.initialSettings as Record<string, unknown>)
+		this.sendToExtension({ type: "updateSettings", updatedSettings: this.initialSettings })
 	}
 
-	private flushPendingMessages(): void {
-		if (this.pendingMessages.length > 0) {
-			for (const message of this.pendingMessages) {
-				this.emit("webviewMessage", message)
-			}
-			this.pendingMessages = []
-		}
+	public isInInitialSetup(): boolean {
+		return !this.isReady
 	}
 
 	// ==========================================================================
 	// Message Handling
 	// ==========================================================================
 
-	sendToExtension(message: WebviewMessage): void {
-		if (!this.isWebviewReady) {
-			this.pendingMessages.push(message)
-			return
-		}
-		this.emit("webviewMessage", message)
-	}
-
-	/**
-	 * Handle incoming messages from extension.
-	 * Forward to client (single source of truth).
-	 */
-	private handleExtensionMessage(msg: ExtensionMessage): void {
-		// Track mode changes
-		if (msg.type === "state" && msg.state?.mode && typeof msg.state.mode === "string") {
-			this.currentMode = msg.state.mode
+	public sendToExtension(message: WebviewMessage): void {
+		if (!this.isReady) {
+			throw new Error("You cannot send messages to the extension before it is ready")
 		}
 
-		// Forward to client - it's the single source of truth
-		this.client.handleMessage(msg)
-
-		// Handle modes separately
-		if (msg.type === "modes") {
-			this.emit("modesUpdated", msg)
-		}
+		this.emit("webviewMessage", message)
 	}
 
 	// ==========================================================================
 	// Task Management
 	// ==========================================================================
 
-	private applyRuntimeSettings(settings: RooCodeSettings): void {
-		const activeMode = this.currentMode || this.options.mode
-		if (activeMode) {
-			settings.mode = activeMode
-		}
-
-		if (this.options.reasoningEffort && this.options.reasoningEffort !== "unspecified") {
-			if (this.options.reasoningEffort === "disabled") {
-				settings.enableReasoningEffort = false
-			} else {
-				settings.enableReasoningEffort = true
-				settings.reasoningEffort = this.options.reasoningEffort
-			}
-		}
-
-		setRuntimeConfigValues("roo-cline", settings as Record<string, unknown>)
-	}
-
-	async runTask(prompt: string): Promise<void> {
-		if (!this.isWebviewReady) {
-			await new Promise<void>((resolve) => this.once("webviewReady", resolve))
-		}
-
-		// Send initial webview messages to trigger proper extension initialization
-		// This is critical for the extension to start sending state updates properly
-		this.sendToExtension({ type: "webviewDidLaunch" })
-
-		const baseSettings: RooCodeSettings = {
-			commandExecutionTimeout: 30,
-			browserToolEnabled: false,
-			enableCheckpoints: false,
-			...getProviderSettings(this.options.provider, this.options.apiKey, this.options.model),
-		}
-
-		const settings: RooCodeSettings = this.options.nonInteractive
-			? {
-					autoApprovalEnabled: true,
-					alwaysAllowReadOnly: true,
-					alwaysAllowReadOnlyOutsideWorkspace: true,
-					alwaysAllowWrite: true,
-					alwaysAllowWriteOutsideWorkspace: true,
-					alwaysAllowWriteProtected: true,
-					alwaysAllowBrowser: true,
-					alwaysAllowMcp: true,
-					alwaysAllowModeSwitch: true,
-					alwaysAllowSubtasks: true,
-					alwaysAllowExecute: true,
-					allowedCommands: ["*"],
-					...baseSettings,
-				}
-			: {
-					autoApprovalEnabled: false,
-					...baseSettings,
-				}
-
-		this.applyRuntimeSettings(settings)
-		this.sendToExtension({ type: "updateSettings", updatedSettings: settings })
-		await new Promise<void>((resolve) => setTimeout(resolve, 100))
+	public async runTask(prompt: string): Promise<void> {
 		this.sendToExtension({ type: "newTask", text: prompt })
-		await this.waitForCompletion()
-	}
 
-	private waitForCompletion(timeoutMs: number = 110000): Promise<void> {
 		return new Promise((resolve, reject) => {
 			let timeoutId: NodeJS.Timeout | null = null
+			const timeoutMs: number = 110_000
 
 			const completeHandler = () => {
 				cleanup()
 				resolve()
 			}
-			const errorHandler = (error: string) => {
-				cleanup()
-				reject(new Error(error))
-			}
-			const timeoutHandler = () => {
+
+			const errorHandler = (error: Error) => {
 				cleanup()
-				reject(
-					new Error(`Task completion timeout after ${timeoutMs}ms - no completion or error event received`),
-				)
+				reject(error)
 			}
+
 			const cleanup = () => {
 				if (timeoutId) {
 					clearTimeout(timeoutId)
 					timeoutId = null
 				}
-				this.off("taskComplete", completeHandler)
-				this.off("taskError", errorHandler)
+
+				this.client.off("taskCompleted", completeHandler)
+				this.client.off("error", errorHandler)
 			}
 
-			// Set timeout to prevent indefinite hanging
-			timeoutId = setTimeout(timeoutHandler, timeoutMs)
+			// Set timeout to prevent indefinite hanging.
+			timeoutId = setTimeout(() => {
+				cleanup()
+				reject(
+					new Error(`Task completion timeout after ${timeoutMs}ms - no completion or error event received`),
+				)
+			}, timeoutMs)
 
-			this.once("taskComplete", completeHandler)
-			this.once("taskError", errorHandler)
+			this.client.once("taskCompleted", completeHandler)
+			this.client.once("error", errorHandler)
 		})
 	}
 
 	// ==========================================================================
-	// Public Agent State API (delegated to ExtensionClient)
+	// Public Agent State API
 	// ==========================================================================
 
 	/**
 	 * Get the current agent loop state.
 	 */
-	getAgentState(): AgentStateInfo {
+	public getAgentState(): AgentStateInfo {
 		return this.client.getAgentState()
 	}
 
 	/**
 	 * Check if the agent is currently waiting for user input.
 	 */
-	isWaitingForInput(): boolean {
+	public isWaitingForInput(): boolean {
 		return this.client.getAgentState().isWaitingForInput
 	}
 
-	/**
-	 * Check if the agent is currently running.
-	 */
-	isAgentRunning(): boolean {
-		return this.client.getAgentState().isRunning
-	}
-
-	/**
-	 * Get the current agent loop state enum value.
-	 */
-	getAgentLoopState(): AgentLoopState {
-		return this.client.getAgentState().state
-	}
-
-	/**
-	 * Get the underlying ExtensionClient for advanced use cases.
-	 */
-	getExtensionClient(): ExtensionClient {
-		return this.client
-	}
-
-	/**
-	 * Get the OutputManager for advanced output control.
-	 */
-	getOutputManager(): OutputManager {
-		return this.outputManager
-	}
-
-	/**
-	 * Get the PromptManager for advanced prompting.
-	 */
-	getPromptManager(): PromptManager {
-		return this.promptManager
-	}
-
-	/**
-	 * Get the AskDispatcher for advanced ask handling.
-	 */
-	getAskDispatcher(): AskDispatcher {
-		return this.askDispatcher
-	}
-
 	// ==========================================================================
 	// Cleanup
 	// ==========================================================================
@@ -644,7 +521,6 @@ export class ExtensionHost extends EventEmitter {
 		this.vscode = null
 		this.extensionModule = null
 		this.extensionAPI = null
-		this.webviewProviders.clear()
 
 		// Clear globals.
 		delete (global as Record<string, unknown>).vscode

apps/cli/src/agent/message-processor.ts

@@ -161,7 +161,21 @@ export class MessageProcessor {
 			return
 		}
 
-		const { clineMessages } = message.state
+		const { clineMessages, mode } = message.state
+
+		// Track mode changes.
+		if (mode && typeof mode === "string") {
+			const previousMode = this.store.getCurrentMode()
+
+			if (previousMode !== mode) {
+				if (this.options.debug) {
+					debugLog("[MessageProcessor] Mode changed", { from: previousMode, to: mode })
+				}
+
+				this.store.setCurrentMode(mode)
+				this.emitter.emit("modeChanged", { previousMode, currentMode: mode })
+			}
+		}
 
 		if (!clineMessages) {
 			if (this.options.debug) {
@@ -170,7 +184,7 @@ export class MessageProcessor {
 			return
 		}
 
-		// Get previous state for comparison
+		// Get previous state for comparison.
 		const previousState = this.store.getAgentState()
 
 		// Update the store with new messages

apps/cli/src/agent/state-store.ts

@@ -48,6 +48,12 @@ export interface StoreState {
 	 */
 	lastUpdatedAt: number
 
+	/**
+	 * The current mode (e.g., "code", "architect", "ask").
+	 * Tracked from state messages received from the extension.
+	 */
+	currentMode: string | undefined
+
 	/**
 	 * Optional: Cache of extension state fields we might need.
 	 * This is a subset of the full ExtensionState.
@@ -64,6 +70,7 @@ function createInitialState(): StoreState {
 		agentState: detectAgentState([]),
 		isInitialized: false,
 		lastUpdatedAt: Date.now(),
+		currentMode: undefined,
 	}
 }
 
@@ -183,6 +190,13 @@ export class StateStore {
 		return this.state.agentState.state
 	}
 
+	/**
+	 * Get the current mode (e.g., "code", "architect", "ask").
+	 */
+	getCurrentMode(): string | undefined {
+		return this.state.currentMode
+	}
+
 	// ===========================================================================
 	// State Updates
 	// ===========================================================================
@@ -203,6 +217,7 @@ export class StateStore {
 			agentState: newAgentState,
 			isInitialized: true,
 			lastUpdatedAt: Date.now(),
+			currentMode: this.state.currentMode, // Preserve mode across message updates
 		})
 
 		return previousAgentState
@@ -249,10 +264,27 @@ export class StateStore {
 			agentState: detectAgentState([]),
 			isInitialized: true, // Still initialized, just empty
 			lastUpdatedAt: Date.now(),
+			currentMode: this.state.currentMode, // Preserve mode when clearing task
 			extensionState: undefined,
 		})
 	}
 
+	/**
+	 * Set the current mode.
+	 * Called when mode changes are detected from extension state messages.
+	 *
+	 * @param mode - The new mode value
+	 */
+	setCurrentMode(mode: string | undefined): void {
+		if (this.state.currentMode !== mode) {
+			this.updateState({
+				...this.state,
+				currentMode: mode,
+				lastUpdatedAt: Date.now(),
+			})
+		}
+	}
+
 	/**
 	 * Reset to completely uninitialized state.
 	 * Called on disconnect or reset.
@@ -366,6 +398,7 @@ export function getDefaultStore(): StateStore {
 	if (!defaultStore) {
 		defaultStore = new StateStore()
 	}
+
 	return defaultStore
 }
 
@@ -377,5 +410,6 @@ export function resetDefaultStore(): void {
 	if (defaultStore) {
 		defaultStore.reset()
 	}
+
 	defaultStore = null
 }

apps/cli/src/lib/storage/ephemeral.ts

@@ -0,0 +1,10 @@
+import path from "path"
+import os from "os"
+import fs from "fs"
+
+export async function createEphemeralStorageDir(): Promise<string> {
+	const uniqueId = `${Date.now()}-${Math.random().toString(36).slice(2, 11)}`
+	const tmpDir = path.join(os.tmpdir(), `roo-cli-${uniqueId}`)
+	await fs.promises.mkdir(tmpDir, { recursive: true })
+	return tmpDir
+}

apps/cli/src/lib/storage/index.ts

@@ -1,3 +1,4 @@
 export * from "./config-dir.js"
 export * from "./settings.js"
 export * from "./credentials.js"
+export * from "./ephemeral.js"

apps/cli/src/ui/App.tsx

@@ -1,9 +1,8 @@
 import { Box, Text, useApp, useInput } from "ink"
 import { Select } from "@inkjs/ui"
 import { useState, useEffect, useCallback, useRef, useMemo } from "react"
-import type { WebviewMessage } from "@roo-code/types"
 
-import { ExtensionHostOptions } from "@/agent/index.js"
+import { ExtensionHostInterface, ExtensionHostOptions } from "@/agent/index.js"
 
 import { getGlobalCommandsForAutocomplete } from "@/lib/utils/commands.js"
 import { arePathsEqual } from "@/lib/utils/path.js"
@@ -59,15 +58,6 @@ import ScrollIndicator from "./components/ScrollIndicator.js"
 
 const PICKER_HEIGHT = 10
 
-interface ExtensionHostInterface {
-	// eslint-disable-next-line @typescript-eslint/no-explicit-any
-	on(event: string, handler: (...args: any[]) => void): void
-	activate(): Promise<void>
-	runTask(prompt: string): Promise<void>
-	sendToExtension(message: WebviewMessage): void
-	dispose(): Promise<void>
-}
-
 export interface TUIAppProps extends ExtensionHostOptions {
 	initialPrompt: string
 	debug: boolean

apps/cli/src/ui/__tests__/store.test.ts

@@ -184,47 +184,47 @@ describe("useCLIStore", () => {
 		it("should support the full task resumption workflow", () => {
 			const store = useCLIStore.getState
 
-			// Step 1: Initial state with task history and modes from webviewDidLaunch
+			// Step 1: Initial state with task history and modes from webviewDidLaunch.
 			store().setTaskHistory([{ id: "task1", task: "Previous task", workspace: "/test", ts: Date.now() }])
 			store().setAvailableModes([{ key: "code", slug: "code", name: "Code" }])
 			store().setAllSlashCommands([{ key: "new", name: "new", source: "global" as const }])
 
-			// Step 2: User starts a new task
+			// Step 2: User starts a new task.
 			store().setHasStartedTask(true)
 			store().addMessage({ id: "1", role: "user", content: "New task" })
 			store().addMessage({ id: "2", role: "assistant", content: "Working on it..." })
 			store().setLoading(true)
 
-			// Verify current state
+			// Verify current state.
 			expect(store().messages.length).toBe(2)
 			expect(store().hasStartedTask).toBe(true)
 
-			// Step 3: User selects a task from history to resume
-			// This triggers resetForTaskSwitch + setIsResumingTask(true)
+			// Step 3: User selects a task from history to resume.
+			// This triggers resetForTaskSwitch + setIsResumingTask(true).
 			store().resetForTaskSwitch()
 			store().setIsResumingTask(true)
 
-			// Verify task-specific state is cleared but global state preserved
+			// Verify task-specific state is cleared but global state preserved.
 			expect(store().messages).toEqual([])
 			expect(store().isLoading).toBe(false)
 			expect(store().hasStartedTask).toBe(false)
-			expect(store().isResumingTask).toBe(true) // Flag is set
-			expect(store().taskHistory.length).toBe(1) // Preserved
-			expect(store().availableModes.length).toBe(1) // Preserved
-			expect(store().allSlashCommands.length).toBe(1) // Preserved
+			expect(store().isResumingTask).toBe(true) // Flag is set.
+			expect(store().taskHistory.length).toBe(1) // Preserved.
+			expect(store().availableModes.length).toBe(1) // Preserved.
+			expect(store().allSlashCommands.length).toBe(1) // Preserved.
 
 			// Step 4: Extension sends state message with clineMessages
-			// (simulated by adding messages)
+			// (simulated by adding messages).
 			store().addMessage({ id: "old1", role: "user", content: "Previous task prompt" })
 			store().addMessage({ id: "old2", role: "assistant", content: "Previous response" })
 
-			// Step 5: After processing state, isResumingTask should be cleared
+			// Step 5: After processing state, isResumingTask should be cleared.
 			store().setIsResumingTask(false)
 
-			// Final verification
+			// Final verification.
 			expect(store().isResumingTask).toBe(false)
 			expect(store().messages.length).toBe(2)
-			expect(store().taskHistory.length).toBe(1) // Still preserved
+			expect(store().taskHistory.length).toBe(1) // Still preserved.
 		})
 
 		it("should allow reading isResumingTask synchronously during message processing", () => {

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

@@ -3,19 +3,10 @@ import { useApp } from "ink"
 import { randomUUID } from "crypto"
 import type { ExtensionMessage, WebviewMessage } from "@roo-code/types"
 
-import { ExtensionHostOptions } from "@/agent/index.js"
+import { ExtensionHostInterface, ExtensionHostOptions } from "@/agent/index.js"
 
 import { useCLIStore } from "../store.js"
 
-interface ExtensionHostInterface {
-	// eslint-disable-next-line @typescript-eslint/no-explicit-any
-	on(event: string, handler: (...args: any[]) => void): void
-	activate(): Promise<void>
-	runTask(prompt: string): Promise<void>
-	sendToExtension(message: WebviewMessage): void
-	dispose(): Promise<void>
-}
-
 export interface UseExtensionHostOptions extends ExtensionHostOptions {
 	initialPrompt?: string
 	exitOnComplete?: boolean
@@ -89,9 +80,9 @@ export function useExtensionHost({
 				hostRef.current = host
 				isReadyRef.current = true
 
-				host.on("extensionWebviewMessage", onExtensionMessage)
+				host.on("extensionWebviewMessage", (msg) => onExtensionMessage(msg as ExtensionMessage))
 
-				host.on("taskComplete", async () => {
+				host.client.on("taskCompleted", async () => {
 					setComplete(true)
 					setLoading(false)
 
@@ -102,8 +93,8 @@ export function useExtensionHost({
 					}
 				})
 
-				host.on("taskError", (err: string) => {
-					setError(err)
+				host.client.on("error", (err: Error) => {
+					setError(err.message)
 					setLoading(false)
 				})
 
@@ -111,7 +102,6 @@ export function useExtensionHost({
 
 				// Request initial state from extension (triggers
 				// postStateToWebview which includes taskHistory).
-				host.sendToExtension({ type: "webviewDidLaunch" })
 				host.sendToExtension({ type: "requestCommands" })
 				host.sendToExtension({ type: "requestModes" })
 
@@ -136,28 +126,25 @@ export function useExtensionHost({
 		}
 	}, []) // Run once on mount
 
-	// Stable sendToExtension - uses ref to always access current host
-	// This function reference never changes, preventing downstream useCallback/useMemo invalidations
+	// Stable sendToExtension - uses ref to always access current host.
+	// This function reference never changes, preventing downstream
+	// useCallback/useMemo invalidations.
 	const sendToExtension = useCallback((msg: WebviewMessage) => {
 		hostRef.current?.sendToExtension(msg)
 	}, [])
 
-	// Stable runTask - uses ref to always access current host
+	// Stable runTask - uses ref to always access current host.
 	const runTask = useCallback((prompt: string): Promise<void> => {
 		if (!hostRef.current) {
 			return Promise.reject(new Error("Extension host not ready"))
 		}
+
 		return hostRef.current.runTask(prompt)
 	}, [])
 
-	// Memoized return object to prevent unnecessary re-renders in consumers
+	// Memoized return object to prevent unnecessary re-renders in consumers.
 	return useMemo(
-		() => ({
-			isReady: isReadyRef.current,
-			sendToExtension,
-			runTask,
-			cleanup,
-		}),
+		() => ({ isReady: isReadyRef.current, sendToExtension, runTask, cleanup }),
 		[sendToExtension, runTask, cleanup],
 	)
 }

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

@@ -73,14 +73,15 @@ export function useTaskSubmit({
 					const globalCommand = getGlobalCommand(commandMatch[1])
 
 					if (globalCommand?.action === "clearTask") {
-						// Reset CLI state and send clearTask to extension
+						// Reset CLI state and send clearTask to extension.
 						useCLIStore.getState().reset()
-						// Reset component-level refs to avoid stale message tracking
+
+						// Reset component-level refs to avoid stale message tracking.
 						seenMessageIds.current.clear()
 						firstTextMessageSkipped.current = false
 						sendToExtension({ type: "clearTask" })
-						// Re-request state, commands and modes since reset() cleared them
-						sendToExtension({ type: "webviewDidLaunch" })
+
+						// Re-request state, commands and modes since reset() cleared them.
 						sendToExtension({ type: "requestCommands" })
 						sendToExtension({ type: "requestModes" })
 						return

packages/vscode-shim/src/index.ts

@@ -80,6 +80,9 @@ export {
 	type FileStat,
 	type Terminal,
 	type CancellationToken,
+	type IExtensionHost,
+	type ExtensionHostEventMap,
+	type ExtensionHostEventName,
 } from "./vscode.js"
 
 // Export utilities

packages/vscode-shim/src/interfaces/extension-host.ts

@@ -0,0 +1,89 @@
+/**
+ * Interface defining the contract that an ExtensionHost must implement
+ * to work with the vscode-shim WindowAPI.
+ *
+ * This interface is used implicitly by WindowAPI when accessing global.__extensionHost.
+ * The ExtensionHost implementation (e.g., in apps/cli) must satisfy this contract.
+ */
+
+import type { WebviewViewProvider } from "./webview.js"
+
+/**
+ * Core event map for ExtensionHost communication.
+ * Maps event names to their payload types.
+ *
+ * - "extensionWebviewMessage": Messages from the extension to the webview/CLI
+ * - "webviewMessage": Messages from the webview/CLI to the extension
+ */
+export interface ExtensionHostEventMap {
+	extensionWebviewMessage: unknown
+	webviewMessage: unknown
+}
+
+/**
+ * Allowed event names for ExtensionHost communication.
+ */
+export type ExtensionHostEventName = keyof ExtensionHostEventMap
+
+/**
+ * ExtensionHost interface for bridging the vscode-shim with the actual extension host.
+ *
+ * The ExtensionHost acts as a message broker between the extension and the CLI/webview,
+ * providing event-based communication and webview provider registration.
+ *
+ * @template TEventMap - Event map type that must include the core ExtensionHostEventMap events.
+ *                       Implementations can extend this with additional events.
+ */
+export interface IExtensionHost<TEventMap extends ExtensionHostEventMap = ExtensionHostEventMap> {
+	/**
+	 * Register a webview view provider with a specific view ID.
+	 * Called by WindowAPI.registerWebviewViewProvider to allow the extension host
+	 * to track registered providers.
+	 *
+	 * @param viewId - The unique identifier for the webview view
+	 * @param provider - The webview view provider to register
+	 */
+	registerWebviewProvider(viewId: string, provider: WebviewViewProvider): void
+
+	/**
+	 * Unregister a previously registered webview view provider.
+	 * Called when disposing of a webview registration.
+	 *
+	 * @param viewId - The unique identifier of the webview view to unregister
+	 */
+	unregisterWebviewProvider(viewId: string): void
+
+	/**
+	 * Check if the extension host is in its initial setup phase.
+	 * Used to determine if certain actions should be deferred until setup completes.
+	 *
+	 * @returns true if initial setup is in progress, false otherwise
+	 */
+	isInInitialSetup(): boolean
+
+	/**
+	 * Mark the webview as ready, signaling that initial setup has completed.
+	 * This should be called after resolveWebviewView completes successfully.
+	 */
+	markWebviewReady(): void
+
+	/**
+	 * Emit an event to registered listeners.
+	 * Used for forwarding messages from the extension to the webview/CLI.
+	 *
+	 * @param event - The event name to emit
+	 * @param message - The message payload to send with the event
+	 * @returns true if the event had listeners, false otherwise
+	 */
+	emit<K extends keyof TEventMap>(event: K, message: TEventMap[K]): boolean
+
+	/**
+	 * Register a listener for an event.
+	 * Used for receiving messages from the webview/CLI to the extension.
+	 *
+	 * @param event - The event name to listen for
+	 * @param listener - The callback function to invoke when the event is emitted
+	 * @returns The ExtensionHost instance for chaining
+	 */
+	on<K extends keyof TEventMap>(event: K, listener: (message: TEventMap[K]) => void): this
+}

packages/vscode-shim/src/vscode.ts

@@ -129,6 +129,9 @@ export type {
 	UriHandler,
 } from "./interfaces/webview.js"
 
+// Extension host interface
+export type { IExtensionHost, ExtensionHostEventMap, ExtensionHostEventName } from "./interfaces/extension-host.js"
+
 // Workspace interfaces
 export type {
 	WorkspaceConfiguration,

pnpm-lock.yaml

@@ -109,6 +109,9 @@ importers:
       ink:
         specifier: ^6.6.0
         version: 6.6.0(@types/[email protected])([email protected])
+      p-wait-for:
+        specifier: ^5.0.2
+        version: 5.0.2
       react:
         specifier: ^19.1.0
         version: 19.2.3