From 569055ea9a4b5ef5429b3facd6f79b907d271d3c Mon Sep 17 00:00:00 2001 From: cte Date: Sat, 10 Jan 2026 01:34:24 -0800 Subject: [PATCH] Better type safety --- .../agent/__tests__/extension-host.test.ts | 59 ++++++++++++------- apps/cli/src/agent/extension-host.ts | 39 +++++------- apps/cli/src/ui/hooks/useExtensionHost.ts | 8 +-- packages/vscode-shim/src/index.ts | 2 + .../src/interfaces/extension-host.ts | 26 +++++++- packages/vscode-shim/src/vscode.ts | 2 +- 6 files changed, 82 insertions(+), 54 deletions(-) diff --git a/apps/cli/src/agent/__tests__/extension-host.test.ts b/apps/cli/src/agent/__tests__/extension-host.test.ts index d9cdb11efb..38edf50d28 100644 --- a/apps/cli/src/agent/__tests__/extension-host.test.ts +++ b/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", () => { diff --git a/apps/cli/src/agent/extension-host.ts b/apps/cli/src/agent/extension-host.ts index f42a6c3aa1..da40c990b2 100644 --- a/apps/cli/src/agent/extension-host.ts +++ b/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 } -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 { + client: ExtensionClient activate(): Promise runTask(prompt: string): Promise sendToExtension(message: WebviewMessage): void dispose(): Promise } -export class ExtensionHost extends EventEmitter implements IExtensionHost { +export class ExtensionHost extends EventEmitter implements ExtensionHostInterface { // Extension lifecycle. private vscode: ReturnType | 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) }) } diff --git a/apps/cli/src/ui/hooks/useExtensionHost.ts b/apps/cli/src/ui/hooks/useExtensionHost.ts index 3494cc66ac..91bdac2bf0 100644 --- a/apps/cli/src/ui/hooks/useExtensionHost.ts +++ b/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) }) diff --git a/packages/vscode-shim/src/index.ts b/packages/vscode-shim/src/index.ts index 52f50138f5..8f40746de7 100644 --- a/packages/vscode-shim/src/index.ts +++ b/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 diff --git a/packages/vscode-shim/src/interfaces/extension-host.ts b/packages/vscode-shim/src/interfaces/extension-host.ts index 1b48327291..f485ee6021 100644 --- a/packages/vscode-shim/src/interfaces/extension-host.ts +++ b/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 { /** * 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(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(event: K, listener: (message: TEventMap[K]) => void): this } diff --git a/packages/vscode-shim/src/vscode.ts b/packages/vscode-shim/src/vscode.ts index 9d6c55255c..a25cd1e8d9 100644 --- a/packages/vscode-shim/src/vscode.ts +++ b/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 {