From a707ab2d9bb4db90b0a68912f947428879b325af Mon Sep 17 00:00:00 2001 From: Duong Pham Ngoc Date: Thu, 2 Jul 2026 15:26:19 +0700 Subject: [PATCH] fix(tool-overrides): register built-in renderers during load --- src/tool-overrides.ts | 64 ++++++--- tests/index-integration.test.ts | 19 +-- tests/reload-behavior.test.ts | 155 ++++++++++++---------- tests/tool-overrides-registration.test.ts | 35 ++++- 4 files changed, 167 insertions(+), 106 deletions(-) diff --git a/src/tool-overrides.ts b/src/tool-overrides.ts index 88116a9..1622c44 100644 --- a/src/tool-overrides.ts +++ b/src/tool-overrides.ts @@ -1554,6 +1554,43 @@ export function registerToolDisplayOverrides( const registeredBuiltInToolOverrides = new Set(); const deferredBuiltInToolOverrides = new Map void>(); + const getExternalBuiltInToolOwner = ( + toolName: BuiltInToolOverrideName, + ): { toolName: BuiltInToolOverrideName; source: string; path: string } | undefined => { + let allTools: unknown[] = []; + try { + allTools = pi.getAllTools(); + } catch (error) { + logToolDisplayDebug("Built-in tool override ownership discovery unavailable during load.", error); + return undefined; + } + + const currentOwner = allTools.find((tool) => getTextField(tool, "name") === toolName); + const sourceInfo = toRecord(toRecord(currentOwner).sourceInfo); + const source = getTextField(sourceInfo, "source"); + if (!currentOwner || !source || source === "builtin") { + return undefined; + } + + return { + toolName, + source, + path: getTextField(sourceInfo, "path") ?? "unknown", + }; + }; + + const shouldSkipBuiltInToolOverride = ( + toolName: BuiltInToolOverrideName, + ): boolean => { + const externalOwner = getExternalBuiltInToolOwner(toolName); + if (!externalOwner) { + return false; + } + + logToolDisplayDebug("Skipped built-in tool display override because another tool owner is active.", externalOwner); + return true; + }; + const registerIfOwned = ( toolName: BuiltInToolOverrideName, register: () => void, @@ -1573,7 +1610,9 @@ export function registerToolDisplayOverrides( if (options.deferUntilBuiltinOwner) { deferredBuiltInToolOverrides.set(toolName as DeferredBuiltInToolOverrideName, registerOnce); - return; + if (shouldSkipBuiltInToolOverride(toolName)) { + return; + } } registerOnce(); @@ -1584,34 +1623,15 @@ export function registerToolDisplayOverrides( return; } - let allTools: unknown[] = []; - try { - allTools = pi.getAllTools(); - } catch (error) { - logToolDisplayDebug("Built-in tool override ownership discovery failed.", error); - return; - } - for (const [toolName, register] of deferredBuiltInToolOverrides) { if ( registeredBuiltInToolOverrides.has(toolName) || - !getConfig().registerToolOverrides[toolName] + !getConfig().registerToolOverrides[toolName] || + shouldSkipBuiltInToolOverride(toolName) ) { continue; } - const currentOwner = allTools.find((tool) => getTextField(tool, "name") === toolName); - const sourceInfo = toRecord(toRecord(currentOwner).sourceInfo); - const source = getTextField(sourceInfo, "source"); - if (currentOwner && source && source !== "builtin") { - logToolDisplayDebug("Skipped built-in tool display override because another tool owner is active.", { - toolName, - source, - path: getTextField(sourceInfo, "path") ?? "unknown", - }); - continue; - } - register(); } }; diff --git a/tests/index-integration.test.ts b/tests/index-integration.test.ts index b3917d7..6a78971 100644 --- a/tests/index-integration.test.ts +++ b/tests/index-integration.test.ts @@ -98,15 +98,12 @@ test("entry point registers built-in tool overrides", () => { const { api, capturedTools } = createApiStub(); toolDisplayExtension(api); - const toolNames = capturedTools.map((t) => t.name); - // find, ls, write are registered immediately; read/grep/edit/bash are deferred - assert.ok(toolNames.includes("find"), "find tool override registered"); - assert.ok(toolNames.includes("ls"), "ls tool override registered"); - assert.ok(toolNames.includes("write"), "write tool override registered"); - - // Disabled tools (if config disables them) would not appear; the default - // config enables all, so we expect at least these 3 immediately. - assert.ok(toolNames.length >= 3, "at least 3 tool overrides registered immediately"); + const toolNames = capturedTools.map((t) => t.name).sort(); + assert.deepEqual( + toolNames, + ["bash", "edit", "find", "grep", "ls", "read", "write"], + "built-in renderer overrides are registered during extension load", + ); }); test("session_start handler refreshes capabilities and notifies pending errors", async () => { @@ -340,9 +337,7 @@ test("overridden tools preserve promptSnippet and promptGuidelines from built-in const byName = new Map(capturedTools.map((t) => [t.name, t])); - // read (deferred) won't be registered immediately; it's deferred - // So we only check tools registered immediately - for (const name of ["find", "ls", "write"] as const) { + for (const name of ["read", "grep", "find", "ls", "bash", "edit", "write"] as const) { const tool = byName.get(name); assert.ok(tool, `${name} is registered`); // promptSnippet should be a non-empty string or undefined diff --git a/tests/reload-behavior.test.ts b/tests/reload-behavior.test.ts index 7699724..6aa3114 100644 --- a/tests/reload-behavior.test.ts +++ b/tests/reload-behavior.test.ts @@ -1,4 +1,7 @@ import assert from "node:assert/strict"; +import { mkdtempSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; import test from "node:test"; import { UserMessageComponent, @@ -85,7 +88,7 @@ function createApiStub( /** * Create a stub for registerToolDisplayOverrides tests that need event-driven - * deferred registration (read/edit/grep deferral). + * post-load registration safety checks. */ function createExtensionApiStub(allTools: unknown[] = []): { api: ExtensionAPI; @@ -109,6 +112,15 @@ function createExtensionApiStub(allTools: unknown[] = []): { return { api, registeredTools, eventHandlers }; } +function withTempDir(name: string, run: (dir: string) => void): void { + const dir = mkdtempSync(join(tmpdir(), name)); + try { + run(dir); + } finally { + rmSync(dir, { recursive: true, force: true }); + } +} + /** Minimal theme stub for render calls. */ const stubTheme = { fg: (_color: string, text: string) => text, @@ -156,11 +168,11 @@ test("2: built-in tool overrides are re-registered on reload", () => { toolDisplayExtension(api); const countAfterReload = capturedTools.length; - // Each call to registerToolDisplayOverrides registers the same built-in - // tools again (find, ls, write immediately; read/grep/edit/bash deferred). + // Each call to registerToolDisplayOverrides registers every built-in tool + // immediately so restored history can capture extension renderers. assert.ok( - countAfterReload >= countBeforeReload + 3, - "at least 3 tools re-registered on reload", + countAfterReload >= countBeforeReload + 7, + "all built-in tools re-registered on reload", ); // Verify tool names appear multiple times, meaning they were re-registered @@ -168,18 +180,12 @@ test("2: built-in tool overrides are re-registered on reload", () => { for (const tool of capturedTools) { toolNameCounts.set(tool.name, (toolNameCounts.get(tool.name) ?? 0) + 1); } - assert.ok( - (toolNameCounts.get("find") ?? 0) >= 2, - "find is registered at least twice (two calls)", - ); - assert.ok( - (toolNameCounts.get("ls") ?? 0) >= 2, - "ls is registered at least twice (two calls)", - ); - assert.ok( - (toolNameCounts.get("write") ?? 0) >= 2, - "write is registered at least twice (two calls)", - ); + for (const name of ["bash", "edit", "find", "grep", "ls", "read", "write"] as const) { + assert.ok( + (toolNameCounts.get(name) ?? 0) >= 2, + `${name} is registered at least twice (two calls)`, + ); + } }); test("2: re-registered tools have renderCall and renderResult functions after reload", () => { @@ -195,9 +201,6 @@ test("2: re-registered tools have renderCall and renderResult functions after re if (secondCallTools.length > 0) { for (const tool of secondCallTools) { - if (tool.name === "read" || tool.name === "edit" || tool.name === "grep") { - continue; // Deferred - not registered immediately - } assert.ok( typeof tool.renderCall === "function", `${tool.name} from reload has renderCall`, @@ -210,35 +213,40 @@ test("2: re-registered tools have renderCall and renderResult functions after re } }); -test("2: deferred tool overrides (read/edit/grep) register after before_agent_start on reload", async () => { +test("2: built-in tool overrides register before before_agent_start on reload", async () => { const { api, registeredTools, eventHandlers } = createExtensionApiStub(); + const expectedBuiltIns = ["bash", "edit", "find", "grep", "ls", "read", "write"]; - // First call + // First call registers all built-in renderers immediately so restored + // history can capture them before session_start runs. registerToolDisplayOverrides(api, () => DEFAULT_TOOL_DISPLAY_CONFIG); - const firstImmediate = registeredTools.map((t) => t.name); - - // Deferred tools should not be registered yet - assert.equal(firstImmediate.includes("read"), false); - assert.equal(firstImmediate.includes("edit"), false); - assert.equal(firstImmediate.includes("grep"), false); + assert.deepEqual( + registeredTools.map((t) => t.name).sort(), + expectedBuiltIns, + ); - // Trigger deferred registration + const beforeFirstTrigger = registeredTools.length; await eventHandlers.before_agent_start?.(); - const afterFirstDeferred = registeredTools.map((t) => t.name); - assert.ok(afterFirstDeferred.includes("read"), "read registered after before_agent_start"); - assert.ok(afterFirstDeferred.includes("edit"), "edit registered after before_agent_start"); + assert.equal( + registeredTools.length, + beforeFirstTrigger, + "before_agent_start does not duplicate immediately registered tools", + ); - // Simulate reload: call again + // Simulate reload: call again and verify the second load also registers + // all built-ins before the lifecycle safety pass. registerToolDisplayOverrides(api, () => DEFAULT_TOOL_DISPLAY_CONFIG); - const beforeSecondTrigger = registeredTools.length; + assert.deepEqual( + registeredTools.slice(-expectedBuiltIns.length).map((t) => t.name).sort(), + expectedBuiltIns, + ); - // Trigger deferred registration again + const beforeSecondTrigger = registeredTools.length; await eventHandlers.before_agent_start?.(); - const afterSecondDeferred = registeredTools.length; - - assert.ok( - afterSecondDeferred > beforeSecondTrigger, - "deferred tools re-register on reloaded before_agent_start", + assert.equal( + registeredTools.length, + beforeSecondTrigger, + "reloaded before_agent_start does not duplicate immediately registered tools", ); }); @@ -735,9 +743,6 @@ test("9: calling toolDisplayExtension three times (double reload) is safe", () = // Verify all tool registrations have renderCall/renderResult for (const tool of capturedTools) { - if (tool.name === "read" || tool.name === "edit" || tool.name === "grep") { - continue; // Deferred tools - } if (tool.renderCall !== undefined) { assert.equal( typeof tool.renderCall, @@ -893,7 +898,7 @@ test("11: registerToolDisplayOverrides creates fresh state on each call", () => // Each call to registerToolDisplayOverrides creates new: // - builtInToolCache (cleared) // - registeredBuiltInToolOverrides Set - // - deferredBuiltInToolOverrides Map + // - deferredBuiltInToolOverrides Map used for post-load safety checks // - wrappedMcpToolNames Set // - ToolDisplayApi on globalThis @@ -942,7 +947,7 @@ test("11: each tool override call clones parameters independently", () => { registeredTools.map((t) => [t.name, t.parameters]), ); - // Trigger deferred registration + // Trigger the post-load registration safety pass eventHandlers.before_agent_start?.(); const firstAllParamRefs = new Map( @@ -981,34 +986,42 @@ test("11: each tool override call clones parameters independently", () => { // --------------------------------------------------------------------------- test("12: config-store reloads config on fingerprint change between calls", () => { - const initialResult = loadToolDisplayConfig(); + withTempDir("pi-tool-display-reload-config-", (dir) => { + const configFile = join(dir, "config.json"); + const initialResult = loadToolDisplayConfig(configFile); + + assert.deepEqual( + initialResult.config, + DEFAULT_TOOL_DISPLAY_CONFIG, + "missing isolated config loads defaults", + ); - // Config is cached, but if we change the file, fingerprint changes. - // Since we can't easily change the file in a test, verify that: - // - Loading returns a valid config - // - The config has expected defaults - assert.ok(initialResult.config, "config loaded successfully"); - assert.equal( - initialResult.config.readOutputMode, - DEFAULT_TOOL_DISPLAY_CONFIG.readOutputMode, - ); - assert.equal( - initialResult.config.searchOutputMode, - DEFAULT_TOOL_DISPLAY_CONFIG.searchOutputMode, - ); + const changedConfig: ToolDisplayConfig = { + ...DEFAULT_TOOL_DISPLAY_CONFIG, + readOutputMode: "summary", + searchOutputMode: "count", + }; + writeFileSync(configFile, `${JSON.stringify(changedConfig, null, 2)}\n`, "utf8"); + + const afterFileChangeResult = loadToolDisplayConfig(configFile); + assert.equal( + afterFileChangeResult.config.readOutputMode, + "summary", + "fingerprint change causes a re-read from disk", + ); + assert.equal(afterFileChangeResult.config.searchOutputMode, "count"); - // saveToolDisplayConfig clears the cache, forcing a re-read - const saveResult = saveToolDisplayConfig(initialResult.config); - assert.ok(saveResult.success, "config saved successfully (cache cleared)"); + // saveToolDisplayConfig clears the cache, forcing the next load to re-read. + const saveResult = saveToolDisplayConfig(DEFAULT_TOOL_DISPLAY_CONFIG, configFile); + assert.ok(saveResult.success, "config saved successfully (cache cleared)"); - // After save, cache is cleared. Next load re-reads from disk. - const afterSaveResult = loadToolDisplayConfig(); - assert.ok(afterSaveResult.config, "config re-loaded after cache clear"); - assert.equal( - afterSaveResult.config.readOutputMode, - initialResult.config.readOutputMode, - "re-loaded config matches saved config", - ); + const afterSaveResult = loadToolDisplayConfig(configFile); + assert.deepEqual( + afterSaveResult.config, + DEFAULT_TOOL_DISPLAY_CONFIG, + "re-loaded config matches saved config", + ); + }); }); test("12: extension loads config fresh on each call (no stale cache)", () => { @@ -1207,7 +1220,7 @@ test("lifecycle: full session lifecycle (init→reload→invoke handlers) does n test("reload behavior test suite: summary of all tests", () => { const testNames = [ "1: Basic reload detection", - "2: Tool override restoration (re-registration + deferred)", + "2: Tool override restoration (immediate re-registration)", "3: Bash override cleanup (spinner timer lifecycle)", "4: MCP override cleanup (re-decoration on reload)", "5: User message box cleanup (prototype re-patching)", diff --git a/tests/tool-overrides-registration.test.ts b/tests/tool-overrides-registration.test.ts index c80539e..8bca387 100644 --- a/tests/tool-overrides-registration.test.ts +++ b/tests/tool-overrides-registration.test.ts @@ -88,13 +88,46 @@ function createExtensionApiStub(allTools: unknown[] = []): { return { api, registeredTools, eventHandlers }; } +function createExtensionApiStubWithoutLoadTimeToolDiscovery(): { + api: ExtensionAPI; + registeredTools: RegisteredToolLike[]; + eventHandlers: ToolEventHandlers; +} { + const registeredTools: RegisteredToolLike[] = []; + const eventHandlers: ToolEventHandlers = {}; + const api = { + registerTool(tool: RegisteredToolLike): void { + registeredTools.push(tool); + }, + on(event: keyof ToolEventHandlers, handler: () => Promise | void): void { + eventHandlers[event] = handler; + }, + getAllTools(): unknown[] { + throw new Error("Extension runtime not initialized."); + }, + } as unknown as ExtensionAPI; + + return { api, registeredTools, eventHandlers }; +} + +test("registerToolDisplayOverrides registers built-in renderers before session_start for reload history", () => { + const { api, registeredTools } = createExtensionApiStubWithoutLoadTimeToolDiscovery(); + + registerToolDisplayOverrides(api, () => DEFAULT_TOOL_DISPLAY_CONFIG); + + assert.deepEqual( + registeredTools.map((tool) => tool.name).sort(), + ["bash", "edit", "find", "grep", "ls", "read", "write"], + ); +}); + test("registerToolDisplayOverrides copies built-in prompt metadata onto overridden tools", async () => { const { api, registeredTools, eventHandlers } = createExtensionApiStub(); registerToolDisplayOverrides(api, () => DEFAULT_TOOL_DISPLAY_CONFIG); assert.deepEqual( registeredTools.map((tool) => tool.name).sort(), - ["find", "ls", "write"], + ["bash", "edit", "find", "grep", "ls", "read", "write"], ); await eventHandlers.before_agent_start?.();