fix(tool-overrides): register built-in renderers during load

src/tool-overrides.ts

@@ -1554,6 +1554,43 @@ export function registerToolDisplayOverrides(
   const registeredBuiltInToolOverrides = new Set<BuiltInToolOverrideName>();
   const deferredBuiltInToolOverrides = new Map<DeferredBuiltInToolOverrideName, () => 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();
     }
   };

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

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,30 +168,24 @@ 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
   const toolNameCounts = new Map<string, number>();
   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)",