Better type safety

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

@@ -3,10 +3,11 @@
 import { EventEmitter } from "events"
 import fs from "fs"
 
-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(() => ({
@@ -248,11 +249,14 @@ describe("ExtensionHost", () => {
 			const host = createTestHost()
 			const client = getPrivate(host, "client") as ExtensionClient
 
-			// Simulate extension message
-			host.emit("extensionWebviewMessage", { type: "state", state: { clineMessages: [] } })
+			// Simulate extension message.
+			host.emit("extensionWebviewMessage", {
+				type: "state",
+				state: { clineMessages: [] },
+			} as unknown as ExtensionMessage)
 
-			// 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
+			// 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")
 		})
 	})
@@ -420,40 +424,53 @@ describe("ExtensionHost", () => {
 			host.markWebviewReady()
 
 			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
-			setTimeout(() => host.emit("taskComplete"), 10)
+			// 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)
 
 			await taskPromise
 
 			expect(emitSpy).toHaveBeenCalledWith("webviewMessage", { type: "newTask", text: "test prompt" })
 		})
 
-		it("should resolve when taskComplete is emitted", async () => {
+		it("should resolve when taskCompleted is emitted on client", async () => {
 			const host = createTestHost()
 			host.markWebviewReady()
 
+			const client = getPrivate(host, "client") as ExtensionClient
 			const taskPromise = host.runTask("test prompt")
 
-			// Emit completion after a short delay
-			setTimeout(() => host.emit("taskComplete"), 10)
+			// 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(taskPromise).resolves.toBeUndefined()
 		})
-
-		it("should reject when taskError is emitted", async () => {
-			const host = createTestHost()
-			host.markWebviewReady()
-
-			const taskPromise = host.runTask("test prompt")
-
-			setTimeout(() => host.emit("taskError", "Test error"), 10)
-
-			await expect(taskPromise).rejects.toThrow("Test error")
-		})
 	})
 
 	describe("initial settings", () => {

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

@@ -8,11 +8,12 @@
  * 4. Wiring up managers for output, prompting, and ask handling
  */
 
-import { EventEmitter } from "events"
 import { createRequire } from "module"
 import path from "path"
 import { fileURLToPath } from "url"
 import fs from "fs"
+import { EventEmitter } from "events"
+
 import pWaitFor from "p-wait-for"
 
 import type {
@@ -22,7 +23,7 @@ import type {
 	RooCodeSettings,
 	WebviewMessage,
 } from "@roo-code/types"
-import { createVSCodeAPI, IExtensionHost, 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"
@@ -30,7 +31,7 @@ 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 { WaitingForInputEvent, TaskCompletedEvent } from "./events.js"
 import type { AgentStateInfo } from "./agent-state.js"
 import { ExtensionClient } from "./extension-client.js"
 import { OutputManager } from "./output-manager.js"
@@ -83,16 +84,15 @@ interface WebviewViewProvider {
 	resolveWebviewView?(webviewView: unknown, context: unknown, token: unknown): void | Promise<void>
 }
 
-export interface ExtensionHostInterface {
-	// eslint-disable-next-line @typescript-eslint/no-explicit-any
-	on(event: string, handler: (...args: any[]) => void): void
+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 implements IExtensionHost {
+export class ExtensionHost extends EventEmitter implements ExtensionHostInterface {
 	// Extension lifecycle.
 	private vscode: ReturnType<typeof createVSCodeAPI> | null = null
 	private extensionModule: ExtensionModule | null = null
@@ -124,7 +124,7 @@ export class ExtensionHost extends EventEmitter implements IExtensionHost {
 	 * 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.
@@ -234,11 +234,6 @@ export class ExtensionHost extends EventEmitter implements IExtensionHost {
 	 * 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")
@@ -253,22 +248,16 @@ export class ExtensionHost extends EventEmitter implements IExtensionHost {
 
 		// 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)
-
 			// 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")
 		})
 	}
 
@@ -454,9 +443,9 @@ export class ExtensionHost extends EventEmitter implements IExtensionHost {
 				resolve()
 			}
 
-			const errorHandler = (error: string) => {
+			const errorHandler = (error: Error) => {
 				cleanup()
-				reject(new Error(error))
+				reject(error)
 			}
 
 			const cleanup = () => {
@@ -465,8 +454,8 @@ export class ExtensionHost extends EventEmitter implements IExtensionHost {
 					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.
@@ -477,8 +466,8 @@ export class ExtensionHost extends EventEmitter implements IExtensionHost {
 				)
 			}, timeoutMs)
 
-			this.once("taskComplete", completeHandler)
-			this.once("taskError", errorHandler)
+			this.client.once("taskCompleted", completeHandler)
+			this.client.once("error", errorHandler)
 		})
 	}
 

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

@@ -80,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)
 
@@ -93,8 +93,8 @@ export function useExtensionHost({
 					}
 				})
 
-				host.on("taskError", (err: string) => {
-					setError(err)
+				host.client.on("error", (err: Error) => {
+					setError(err.message)
 					setLoading(false)
 				})
 

packages/vscode-shim/src/index.ts

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

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

@@ -8,13 +8,33 @@
 
 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 {
+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
@@ -55,7 +75,7 @@ export interface IExtensionHost {
 	 * @param message - The message payload to send with the event
 	 * @returns true if the event had listeners, false otherwise
 	 */
-	emit(event: string, message: unknown): boolean
+	emit<K extends keyof TEventMap>(event: K, message: TEventMap[K]): boolean
 
 	/**
 	 * Register a listener for an event.
@@ -65,5 +85,5 @@ export interface IExtensionHost {
 	 * @param listener - The callback function to invoke when the event is emitted
 	 * @returns The ExtensionHost instance for chaining
 	 */
-	on(event: string, listener: (message: unknown) => void): this
+	on<K extends keyof TEventMap>(event: K, listener: (message: TEventMap[K]) => void): this
 }

packages/vscode-shim/src/vscode.ts

@@ -130,7 +130,7 @@ export type {
 } from "./interfaces/webview.js"
 
 // Extension host interface
-export type { IExtensionHost } from "./interfaces/extension-host.js"
+export type { IExtensionHost, ExtensionHostEventMap, ExtensionHostEventName } from "./interfaces/extension-host.js"
 
 // Workspace interfaces
 export type {