From 36fe36b081e84f4c8e9b097c992fd6707dbc840a Mon Sep 17 00:00:00 2001 From: kaiwishc Date: Sun, 2 Aug 2026 23:16:34 +0800 Subject: [PATCH] feat: support factory-based extension widgets --- lib/rpc-manager-widgets.test.mjs | 269 +++++++++++++++++++++++++++++++ lib/rpc-manager.ts | 233 ++++++++++++++++++++++++-- 2 files changed, 492 insertions(+), 10 deletions(-) create mode 100644 lib/rpc-manager-widgets.test.mjs diff --git a/lib/rpc-manager-widgets.test.mjs b/lib/rpc-manager-widgets.test.mjs new file mode 100644 index 00000000..5edeaae1 --- /dev/null +++ b/lib/rpc-manager-widgets.test.mjs @@ -0,0 +1,269 @@ +import assert from "node:assert/strict"; +import test from "node:test"; +import { createJiti } from "jiti"; + +const jiti = createJiti(import.meta.url, { + interopDefault: true, + moduleCache: false, +}); +const { AgentSessionWrapper } = await jiti.import("./rpc-manager.ts"); + +function makeInner(overrides = {}) { + return { + sessionId: "widget-test-session", + sessionFile: undefined, + isStreaming: false, + isCompacting: false, + isBashRunning: false, + autoCompactionEnabled: true, + autoRetryEnabled: true, + model: undefined, + modelRuntime: { + getModel: () => undefined, + refresh: async () => {}, + }, + sessionManager: { getCwd: () => process.cwd() }, + settingsManager: { setProjectTrusted: () => {} }, + agent: { state: {} }, + extensionRunner: { + getRegisteredCommands: () => [], + setUIContext: () => {}, + emit: async () => {}, + }, + promptTemplates: [], + resourceLoader: { getSkills: () => ({ skills: [] }) }, + subscribe: () => () => {}, + getContextUsage: () => null, + getSteeringMessages: () => [], + getFollowUpMessages: () => [], + pendingMessageCount: 0, + dispose: () => {}, + reload: async () => {}, + ...overrides, + }; +} + +function createContext(overrides = {}) { + const events = []; + const wrapper = new AgentSessionWrapper(makeInner(overrides)); + wrapper.onEvent((event) => events.push(event)); + const context = wrapper.createExtensionUiContext(); + return { context, events, wrapper }; +} + +function widgetEvents(events) { + return events.filter((event) => event.type === "extension_ui_request" && event.method === "setWidget"); +} + +function snapshot(wrapper) { + return wrapper.send({ type: "get_state" }).then((state) => state.extensionWidgets); +} + +test("keeps array widgets compatible while maintaining the state snapshot", async () => { + const { context, events, wrapper } = createContext(); + + context.setWidget("array", ["one"], { placement: "belowEditor" }); + assert.deepEqual(widgetEvents(events).at(-1), { + type: "extension_ui_request", + id: widgetEvents(events).at(-1).id, + method: "setWidget", + widgetKey: "array", + widgetLines: ["one"], + widgetPlacement: "belowEditor", + }); + assert.deepEqual(await snapshot(wrapper), [{ key: "array", lines: ["one"], placement: "belowEditor" }]); + + const eventCount = widgetEvents(events).length; + context.setWidget("array", ["two"]); + assert.equal(widgetEvents(events).length, eventCount + 1); + assert.equal(widgetEvents(events).at(-1).widgetLines[0], "two"); + assert.equal(widgetEvents(events).at(-1).widgetPlacement, undefined); + assert.deepEqual(await snapshot(wrapper), [{ key: "array", lines: ["two"], placement: "aboveEditor" }]); + wrapper.destroy(); +}); + +test("renders a factory widget immediately with the fixed headless facade", async () => { + const { context, events, wrapper } = createContext(); + let tui; + let theme; + const widths = []; + + context.setWidget("factory", (receivedTui, receivedTheme) => { + tui = receivedTui; + theme = receivedTheme; + return { + render(width) { + widths.push(width); + return [`${width}:${tui.terminal.columns}`]; + }, + }; + }, { placement: "belowEditor" }); + + const event = widgetEvents(events).at(-1); + assert.equal(tui.terminal.columns, 92); + assert.equal(tui.terminal.kittyProtocolActive, false); + assert.ok(theme); + assert.deepEqual(widths, [92]); + assert.deepEqual(event.widgetLines, ["92:92"]); + assert.equal(event.widgetPlacement, "belowEditor"); + assert.deepEqual(await snapshot(wrapper), [{ key: "factory", lines: ["92:92"], placement: "belowEditor" }]); + wrapper.destroy(); +}); + +test("requestRender replaces the browser event and server snapshot", async () => { + const { context, events, wrapper } = createContext(); + let tui; + let value = "first"; + context.setWidget("refreshable", (receivedTui) => { + tui = receivedTui; + return { render: () => [value] }; + }); + + const before = widgetEvents(events).length; + value = "second"; + tui.requestRender(); + assert.equal(widgetEvents(events).length, before + 1); + assert.deepEqual(widgetEvents(events).at(-1).widgetLines, ["second"]); + assert.deepEqual(await snapshot(wrapper), [{ key: "refreshable", lines: ["second"], placement: "aboveEditor" }]); + wrapper.destroy(); +}); + +test("clearing and replacing widgets disposes components and invalidates old callbacks", () => { + const { context, events, wrapper } = createContext(); + let oldTui; + let oldDisposed = 0; + context.setWidget("shared", (tui) => { + oldTui = tui; + return { render: () => ["old"], dispose: () => { oldDisposed += 1; } }; + }); + const beforeReplace = widgetEvents(events).length; + + context.setWidget("shared", ["array"]); + assert.equal(oldDisposed, 1); + assert.deepEqual(widgetEvents(events).at(-1).widgetLines, ["array"]); + const afterReplace = widgetEvents(events).length; + oldTui.requestRender(); + assert.equal(widgetEvents(events).length, afterReplace); + + let newTui; + let newDisposed = 0; + context.setWidget("shared", (tui) => { + newTui = tui; + return { render: () => ["new"], dispose: () => { newDisposed += 1; } }; + }); + const afterFactoryReplace = widgetEvents(events).length; + oldTui.requestRender(); + assert.equal(widgetEvents(events).length, afterFactoryReplace); + + context.setWidget("shared", undefined); + assert.equal(newDisposed, 1); + const afterClear = widgetEvents(events).length; + newTui.requestRender(); + assert.equal(widgetEvents(events).length, afterClear); + assert.equal(widgetEvents(events).at(-1).widgetLines, undefined); + assert.ok(beforeReplace < afterReplace); + wrapper.destroy(); +}); + +test("isolates factory and render failures with a clear and extension_error", async () => { + const { context, events, wrapper } = createContext(); + context.setWidget("factory-error", () => { + throw new Error("factory failed"); + }); + assert.equal(widgetEvents(events).at(-1).widgetLines, undefined); + assert.equal(events.at(-1).type, "extension_error"); + assert.match(events.at(-1).error, /factory failed/); + assert.deepEqual(await snapshot(wrapper), []); + + context.setWidget("factory-error", () => ({ render: () => ["recovered"] })); + assert.deepEqual(await snapshot(wrapper), [{ key: "factory-error", lines: ["recovered"], placement: "aboveEditor" }]); + context.setWidget("factory-error", undefined); + + context.setWidget("invalid-component", () => ({ + dispose: () => { throw new Error("dispose failed"); }, + })); + assert.equal(widgetEvents(events).at(-1).widgetLines, undefined); + assert.match(events.at(-1).error, /render\(width\)/); + assert.deepEqual(await snapshot(wrapper), []); + + let disposed = 0; + context.setWidget("render-error", () => ({ + render: () => { throw new Error("render failed"); }, + dispose: () => { disposed += 1; }, + })); + assert.equal(disposed, 1); + assert.equal(widgetEvents(events).at(-1).widgetLines, undefined); + assert.match(events.at(-1).error, /render failed/); + assert.deepEqual(await snapshot(wrapper), []); + + context.setWidget("invalid-lines", () => ({ render: () => ["ok", 7] })); + assert.equal(widgetEvents(events).at(-1).widgetLines, undefined); + assert.match(events.at(-1).error, /string\[\]/); + assert.deepEqual(await snapshot(wrapper), []); + + let tui; + let shouldThrow = false; + context.setWidget("late-render", (receivedTui) => { + tui = receivedTui; + return { + render: () => { + if (shouldThrow) throw new Error("late render failed"); + return ["live"]; + }, + }; + }); + shouldThrow = true; + tui.requestRender(); + assert.equal(widgetEvents(events).at(-1).widgetLines, undefined); + assert.match(events.at(-1).error, /late render failed/); + assert.deepEqual(await snapshot(wrapper), []); + wrapper.destroy(); +}); + +test("reload and destroy dispose factory components and invalidate callbacks", async () => { + const reloads = []; + const { context, events, wrapper } = createContext({ + reload: async (options) => { + reloads.push(options); + }, + }); + let tui; + let disposed = 0; + context.setWidget("lifecycle", (receivedTui) => { + tui = receivedTui; + return { render: () => ["live"], dispose: () => { disposed += 1; } }; + }); + + await wrapper.send({ type: "reload" }); + assert.equal(disposed, 1); + assert.equal(reloads.length, 1); + assert.deepEqual(await snapshot(wrapper), []); + const afterReload = widgetEvents(events).length; + tui.requestRender(); + assert.equal(widgetEvents(events).length, afterReload); + + let replacementTui; + context.setWidget("lifecycle", (receivedTui) => { + replacementTui = receivedTui; + return { render: () => ["replacement"], dispose: () => { disposed += 1; } }; + }); + wrapper.destroy(); + assert.equal(disposed, 2); + const afterDestroy = widgetEvents(events).length; + replacementTui.requestRender(); + assert.equal(widgetEvents(events).length, afterDestroy); +}); + +test("command-context reload also clears factory widgets", async () => { + const { context, wrapper } = createContext(); + let disposed = 0; + context.setWidget("command-reload", () => ({ + render: () => ["live"], + dispose: () => { disposed += 1; }, + })); + + await wrapper.createExtensionCommandContextActions().reload(); + assert.equal(disposed, 1); + assert.deepEqual(await snapshot(wrapper), []); + wrapper.destroy(); +}); diff --git a/lib/rpc-manager.ts b/lib/rpc-manager.ts index cfb61100..ebd07961 100644 --- a/lib/rpc-manager.ts +++ b/lib/rpc-manager.ts @@ -13,7 +13,7 @@ import { persistExplicitStartupPreferences } from "./startup-preferences"; import type { SlashCommandInfo } from "@earendil-works/pi-coding-agent"; import type { AgentSessionLike, ExtensionUiContextLike, ToolInfo } from "./pi-types"; import type { ExtensionUiRequest, ExtensionUiResponse, ExtensionWidgetItem } from "./types"; -import { createHeadlessCustomUiTui, DEFAULT_CUSTOM_UI_COLUMNS } from "./custom-ui-terminal"; +import { createHeadlessCustomUiTui, DEFAULT_CUSTOM_UI_COLUMNS, type HeadlessCustomUiTui } from "./custom-ui-terminal"; // ============================================================================ // Types @@ -38,6 +38,22 @@ type CustomUiComponent = { invalidate?: () => void; }; +type ExtensionWidgetComponent = { + render: (width: number) => unknown; + dispose?: () => void; +}; + +type ExtensionWidgetFactory = (tui: HeadlessCustomUiTui, theme: Theme) => unknown; + +type ActiveExtensionWidget = { + key: string; + component: ExtensionWidgetComponent; + placement: "aboveEditor" | "belowEditor"; + generation: number; + clearEmitted: boolean; + rendered: boolean; +}; + type ActiveCustomUi = { component: CustomUiComponent; width: number; @@ -141,6 +157,8 @@ export class AgentSessionWrapper { private activeCustomUis = new Map(); private extensionStatuses = new Map(); private extensionWidgets = new Map(); + private activeExtensionWidgets = new Map(); + private extensionWidgetGenerations = new Map(); private promptRunning = false; private extensionsBound = false; private extensionBindingPromise: Promise | null = null; @@ -586,7 +604,7 @@ export class AgentSessionWrapper { case "reload": { await this.waitForExtensionsBound(); this.extensionStatuses.clear(); - this.extensionWidgets.clear(); + this.resetExtensionWidgetsForReload(); this.syncProjectTrust(); await this.inner.reload(); if (typeof this.inner.bindExtensions !== "function") { @@ -658,6 +676,7 @@ export class AgentSessionWrapper { for (const id of Array.from(this.activeCustomUis.keys())) this.closeCustomUi(id, undefined); this.pendingUiResponses.clear(); this.pendingUiRequests.clear(); + this.clearExtensionWidgets(false); try { this.inner.dispose(); } finally { @@ -705,6 +724,190 @@ export class AgentSessionWrapper { return Array.from(this.extensionWidgets.values()); } + private nextExtensionWidgetGeneration(key: string): number { + const generation = (this.extensionWidgetGenerations.get(key) ?? 0) + 1; + this.extensionWidgetGenerations.set(key, generation); + return generation; + } + + private disposeExtensionWidgetComponent(component: unknown): void { + if (!component || (typeof component !== "object" && typeof component !== "function")) return; + const dispose = (component as { dispose?: unknown }).dispose; + if (typeof dispose !== "function") return; + try { + dispose.call(component); + } catch { + // Ignore dispose errors from extension widgets. + } + } + + private emitExtensionWidgetClear(key: string): void { + this.emit({ + type: "extension_ui_request", + id: randomUUID(), + method: "setWidget", + widgetKey: key, + widgetLines: undefined, + widgetPlacement: undefined, + } as ExtensionUiRequest as AgentEvent); + } + + private clearExtensionWidget(key: string, emitClear = true): number { + const generation = this.nextExtensionWidgetGeneration(key); + + const active = this.activeExtensionWidgets.get(key); + this.activeExtensionWidgets.delete(key); + this.extensionWidgets.delete(key); + if (active) this.disposeExtensionWidgetComponent(active.component); + if (emitClear) this.emitExtensionWidgetClear(key); + return generation; + } + + private clearExtensionWidgets(emitClear: boolean): void { + const keys = new Set([ + ...this.extensionWidgets.keys(), + ...this.activeExtensionWidgets.keys(), + ]); + for (const key of keys) this.clearExtensionWidget(key, emitClear); + } + + private resetExtensionWidgetsForReload(): void { + const factoryKeys = [...this.activeExtensionWidgets.keys()]; + for (const key of factoryKeys) this.clearExtensionWidget(key); + // Keep the existing array-widget reload behavior: snapshots are reset and + // the next extension session_start repopulates them. + this.extensionWidgets.clear(); + } + + private emitExtensionWidgetError(key: string, error: unknown): void { + this.emit({ + type: "extension_error", + extensionPath: `extension-widget:${key}`, + event: "setWidget", + error: error instanceof Error ? error.message : String(error), + }); + } + + private failExtensionWidget( + key: string, + generation: number, + error: unknown, + clearEmitted: boolean, + component?: unknown, + ): void { + if (this.extensionWidgetGenerations.get(key) !== generation) { + this.disposeExtensionWidgetComponent(component); + return; + } + + const active = this.activeExtensionWidgets.get(key); + let shouldEmitClear = !clearEmitted; + if (active?.generation === generation) { + shouldEmitClear = active.rendered || !active.clearEmitted; + this.activeExtensionWidgets.delete(key); + this.disposeExtensionWidgetComponent(active.component); + } else { + this.disposeExtensionWidgetComponent(component); + } + this.extensionWidgets.delete(key); + if (shouldEmitClear) this.emitExtensionWidgetClear(key); + this.emitExtensionWidgetError(key, error); + } + + private renderExtensionWidget(active: ActiveExtensionWidget): void { + if ( + this.activeExtensionWidgets.get(active.key) !== active + || this.extensionWidgetGenerations.get(active.key) !== active.generation + ) return; + + let lines: unknown; + try { + lines = active.component.render(DEFAULT_CUSTOM_UI_COLUMNS); + } catch (error) { + this.failExtensionWidget(active.key, active.generation, error, active.clearEmitted); + return; + } + if (!Array.isArray(lines) || !lines.every((line) => typeof line === "string")) { + this.failExtensionWidget( + active.key, + active.generation, + new Error("Extension widget render must return string[]"), + active.clearEmitted, + ); + return; + } + if ( + this.activeExtensionWidgets.get(active.key) !== active + || this.extensionWidgetGenerations.get(active.key) !== active.generation + ) return; + + const widgetLines = lines as string[]; + this.extensionWidgets.set(active.key, { + key: active.key, + lines: widgetLines, + placement: active.placement, + }); + active.rendered = true; + this.emit({ + type: "extension_ui_request", + id: randomUUID(), + method: "setWidget", + widgetKey: active.key, + widgetLines, + widgetPlacement: active.placement, + } as ExtensionUiRequest as AgentEvent); + } + + private setExtensionWidgetFactory( + key: string, + factory: ExtensionWidgetFactory, + options?: { placement?: "aboveEditor" | "belowEditor" }, + ): void { + const hadPrevious = this.extensionWidgets.has(key) || this.activeExtensionWidgets.has(key); + const generation = this.clearExtensionWidget(key, hadPrevious); + const tui = createHeadlessCustomUiTui(() => { + const active = this.activeExtensionWidgets.get(key); + if (active?.generation === generation) this.renderExtensionWidget(active); + }, DEFAULT_CUSTOM_UI_COLUMNS); + + let component: unknown; + try { + component = factory(tui, PLAIN_TEXT_THEME); + } catch (error) { + this.failExtensionWidget(key, generation, error, hadPrevious); + return; + } + if (this.extensionWidgetGenerations.get(key) !== generation) { + this.disposeExtensionWidgetComponent(component); + return; + } + if ( + !component + || (typeof component !== "object" && typeof component !== "function") + || typeof (component as { render?: unknown }).render !== "function" + ) { + this.failExtensionWidget( + key, + generation, + new Error("Extension widget factory must return a component with render(width)"), + hadPrevious, + component, + ); + return; + } + + const active: ActiveExtensionWidget = { + key, + component: component as ExtensionWidgetComponent, + placement: options?.placement ?? "aboveEditor", + generation, + clearEmitted: hadPrevious, + rendered: false, + }; + this.activeExtensionWidgets.set(key, active); + this.renderExtensionWidget(active); + } + private getCustomUiWidth(options: unknown): number { if (!options || typeof options !== "object") return DEFAULT_CUSTOM_UI_COLUMNS; const overlayOptions = (options as { overlayOptions?: unknown }).overlayOptions; @@ -938,16 +1141,26 @@ export class AgentSessionWrapper { setWorkingIndicator: () => {}, setHiddenThinkingLabel: () => {}, setWidget: (key, content, options) => { + if (typeof content === "function") { + this.setExtensionWidgetFactory( + key, + content as unknown as ExtensionWidgetFactory, + options, + ); + return; + } if (content !== undefined && !Array.isArray(content)) return; if (content === undefined) { - this.extensionWidgets.delete(key); - } else { - this.extensionWidgets.set(key, { - key, - lines: content, - placement: options?.placement ?? "aboveEditor", - }); + this.clearExtensionWidget(key); + return; } + if (this.activeExtensionWidgets.has(key)) this.clearExtensionWidget(key); + else this.nextExtensionWidgetGeneration(key); + this.extensionWidgets.set(key, { + key, + lines: content, + placement: options?.placement ?? "aboveEditor", + }); this.emit({ type: "extension_ui_request", id: randomUUID(), @@ -1012,7 +1225,7 @@ export class AgentSessionWrapper { switchSession: async () => ({ cancelled: true }), reload: async () => { this.extensionStatuses.clear(); - this.extensionWidgets.clear(); + this.resetExtensionWidgetsForReload(); this.syncProjectTrust(); await this.inner.reload({ beforeSessionStart: () => {