fix: decorate late tools before registration

CHANGELOG.md

@@ -7,6 +7,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- Added `preserveCallRenderer` for custom tool overrides that should retain a useful native call header while compacting or hiding only result output.
+
+### Fixed
+- Decorate tools registered by later-loaded extensions before Pi snapshots their definitions.
+- Track decorated tool objects by identity and restore interception state across reloads and defensive double loads.
+
 ## [0.5.0] - 2026-07-03
 
 ### Added

README.md

@@ -197,7 +197,8 @@ Use `customToolOverrides` when another extension registers a noisy top-level too
     "custom_mcp_gateway": {
       "enabled": true,
       "kind": "mcp",
-      "outputMode": "preview"
+      "outputMode": "preview",
+      "preserveCallRenderer": true
     }
   }
 }
@@ -210,6 +211,7 @@ Each entry supports:
 | `enabled` | boolean | `true` | Whether `pi-tool-display` should decorate this custom tool |
 | `kind` | string | `"generic"` | `generic` for plain compact output, or `mcp` for MCP-style call labels and result handling |
 | `outputMode` | string | `"summary"` | `hidden`, `summary`, or `preview` for this custom tool's result output |
+| `preserveCallRenderer` | boolean | `false` | Keep the tool's native call/header renderer while overriding only its result output |
 
 Boolean shorthand is also accepted:
 
@@ -252,7 +254,8 @@ Notes:
     "custom_mcp_gateway": {
       "enabled": true,
       "kind": "mcp",
-      "outputMode": "preview"
+      "outputMode": "preview",
+      "preserveCallRenderer": true
     }
   },
   "enableNativeUserMessageBox": true,

config/config.example.json

@@ -13,7 +13,8 @@
     "ide_find_symbol": {
       "enabled": false,
       "kind": "generic",
-      "outputMode": "summary"
+      "outputMode": "summary",
+      "preserveCallRenderer": true
     },
     "custom_mcp_gateway": {
       "enabled": false,

src/config-store.ts

@@ -176,6 +176,7 @@ export function normalizeCustomToolOverrideEntry(rawEntry: unknown): CustomToolO
 		enabled: toBoolean(source.enabled, true),
 		kind: toCustomToolOverrideKind(source.kind),
 		outputMode: toCustomToolOutputMode(source.outputMode),
+		...(source.preserveCallRenderer === true ? { preserveCallRenderer: true } : {}),
 	};
 }
 

src/index.ts

@@ -33,13 +33,16 @@ function ownershipChanged(
 }
 
 export default function toolDisplayExtension(pi: ExtensionAPI): void {
+  // Pi normally calls session_shutdown before reload, but dispose a prior
+  // generation defensively when loaders invoke this entrypoint twice.
+  disposeAll();
+  resetDisposed();
+
   const initial = loadToolDisplayConfig();
   if (!initial.config.enabled) {
     return;
   }
 
-  resetDisposed();
-
   pi.on("session_shutdown", (event: { reason: string }) => {
     if (event.reason === "reload") {
       disposeAll();

src/tool-overrides.ts

@@ -140,6 +140,7 @@ interface BashToolOverrideOptions {
 const builtInToolCache = new Map<string, BuiltInTools>();
 const RTK_COMPACTION_LABEL = "compacted by RTK";
 export const WRITE_EXECUTION_META_LIMIT = 100;
+const TOOL_DISPLAY_PENDING_DECORATIONS_LIMIT = 100;
 const WRITE_EXECUTION_META_STATE_KEY = "__piToolDisplayWriteExecutionMeta";
 const EDIT_PENDING_PREVIEW_STATE_KEY = "__piToolDisplayEditPendingPreview";
 const WRITE_PENDING_PREVIEW_STATE_KEY = "__piToolDisplayWritePendingPreview";
@@ -166,6 +167,8 @@ export interface ToolDisplayAdapter {
   toolName?: string;
   kind?: ToolDisplayKind;
   overrideExistingRenderers?: boolean;
+  /** Preserve a supplied tool's call renderer while still applying result rendering. */
+  preserveCallRenderer?: boolean;
   pathFields?: string[];
   getPath?: (args: unknown) => string | undefined;
   getEditLineCount?: (args: unknown) => number;
@@ -1478,6 +1481,12 @@ function drainPendingToolDisplayDecorations(api: ToolDisplayApi): void {
     return;
   }
 
+  // Consumer extensions may load before this extension. Keep only the newest
+  // bounded set so a long-lived pre-load queue cannot retain arbitrary tools.
+  if (pendingDecorations.length > TOOL_DISPLAY_PENDING_DECORATIONS_LIMIT) {
+    pendingDecorations.splice(0, pendingDecorations.length - TOOL_DISPLAY_PENDING_DECORATIONS_LIMIT);
+  }
+
   const entries = pendingDecorations.splice(0);
   for (const entry of entries) {
     if (!entry?.tool || typeof entry.tool !== "object") {
@@ -1509,20 +1518,23 @@ function installToolDisplayApi(getConfig: ConfigGetter): ToolDisplayApi {
       const resolvedAdapter = resolveAdapter(tool, adapter);
       const kind = getAdapterKind(tool, resolvedAdapter);
       const overrideExisting = resolvedAdapter.overrideExistingRenderers === true;
+      const preserveCallRenderer = resolvedAdapter.preserveCallRenderer === true;
       const decorated: RuntimeToolDefinition = { ...tool };
 
-      if (resolvedAdapter.renderCall && (overrideExisting || typeof decorated.renderCall !== "function")) {
-        decorated.renderCall = resolvedAdapter.renderCall;
-      } else if (kind === "read" && (overrideExisting || typeof decorated.renderCall !== "function")) {
-        decorated.renderCall = (args: unknown, theme: RenderTheme) => renderReadDisplayCall(args, theme, resolvedAdapter);
-      } else if (kind === "edit" && (overrideExisting || typeof decorated.renderCall !== "function")) {
-        decorated.renderCall = (args: unknown, theme: RenderTheme, context: ToolRenderContextLike) => renderEditDisplayCall(args, theme, context, resolvedAdapter, getConfig);
-      } else if (kind === "mcp" && (overrideExisting || typeof decorated.renderCall !== "function")) {
-        decorated.renderCall = (args: unknown, theme: RenderTheme) => {
-          const toolName = getTextField(decorated, "name") ?? "mcp";
-          const toolLabel = getTextField(decorated, "label") ?? (toolName === "mcp" ? "MCP Proxy" : `MCP ${toolName}`);
-          return formatMcpCallLine(toolName, toolLabel, toRecord(args), theme);
-        };
+      if (!preserveCallRenderer) {
+        if (resolvedAdapter.renderCall && (overrideExisting || typeof decorated.renderCall !== "function")) {
+          decorated.renderCall = resolvedAdapter.renderCall;
+        } else if (kind === "read" && (overrideExisting || typeof decorated.renderCall !== "function")) {
+          decorated.renderCall = (args: unknown, theme: RenderTheme) => renderReadDisplayCall(args, theme, resolvedAdapter);
+        } else if (kind === "edit" && (overrideExisting || typeof decorated.renderCall !== "function")) {
+          decorated.renderCall = (args: unknown, theme: RenderTheme, context: ToolRenderContextLike) => renderEditDisplayCall(args, theme, context, resolvedAdapter, getConfig);
+        } else if (kind === "mcp" && (overrideExisting || typeof decorated.renderCall !== "function")) {
+          decorated.renderCall = (args: unknown, theme: RenderTheme) => {
+            const toolName = getTextField(decorated, "name") ?? "mcp";
+            const toolLabel = getTextField(decorated, "label") ?? (toolName === "mcp" ? "MCP Proxy" : `MCP ${toolName}`);
+            return formatMcpCallLine(toolName, toolLabel, toRecord(args), theme);
+          };
+        }
       }
 
       if (resolvedAdapter.renderResult && (overrideExisting || typeof decorated.renderResult !== "function")) {
@@ -1905,8 +1917,7 @@ export function registerToolDisplayOverrides(
     });
   });
 
-  const wrappedCustomToolNames = new Set<string>();
-  registerCleanup(() => wrappedCustomToolNames.clear());
+  const decoratedCustomTools = new WeakSet<RuntimeToolDefinition>();
 
   const getCustomOverrideForCandidate = (candidate: unknown): {
     toolName: string;
@@ -1927,42 +1938,48 @@ export function registerToolDisplayOverrides(
 
   const decorateCustomToolOverrideCandidate = (candidate: unknown): boolean => {
     const customOverride = getCustomOverrideForCandidate(candidate);
-    if (!customOverride || wrappedCustomToolNames.has(customOverride.toolName)) {
-      return customOverride !== undefined;
+    if (!customOverride) {
+      return false;
     }
 
     const { toolName, override } = customOverride;
     const runtimeTool = candidate as RuntimeToolDefinition;
-    applyToolDisplayDecorationInPlace(
-      runtimeTool,
-      toolDisplayApi,
-      {
-        kind: override.kind,
-        overrideExistingRenderers: true,
-        renderCall(args, theme) {
-          if (override.kind === "mcp") {
-            return formatMcpCallLine("mcp", "MCP Proxy", toRecord(args), theme);
-          }
-          return formatGenericToolCallLine(toolName, args, theme);
-        },
-        renderResult(result, options, theme) {
-          return renderCustomToolResult(
-            result as ToolRenderInput,
-            options,
-            getConfig(),
-            override.outputMode,
-            theme,
-          );
-        },
+    if (decoratedCustomTools.has(runtimeTool)) {
+      return true;
+    }
+
+    const adapter: ToolDisplayAdapter = {
+      kind: override.kind,
+      overrideExistingRenderers: true,
+      preserveCallRenderer: override.preserveCallRenderer === true,
+      renderResult(result, options, theme) {
+        return renderCustomToolResult(
+          result as ToolRenderInput,
+          options,
+          getConfig(),
+          override.outputMode,
+          theme,
+        );
       },
-    );
+    };
+    if (!override.preserveCallRenderer) {
+      adapter.renderCall = (args, theme) => {
+        if (override.kind === "mcp") {
+          return formatMcpCallLine("mcp", "MCP Proxy", toRecord(args), theme);
+        }
+        return formatGenericToolCallLine(toolName, args, theme);
+      };
+    }
 
-    wrappedCustomToolNames.add(toolName);
+    if (!applyToolDisplayDecorationInPlace(runtimeTool, toolDisplayApi, adapter)) {
+      return false;
+    }
+
+    decoratedCustomTools.add(runtimeTool);
     return true;
   };
 
-  const wrappedMcpToolNames = new Set<string>();
-  registerCleanup(() => wrappedMcpToolNames.clear());
+  const decoratedMcpTools = new WeakSet<RuntimeToolDefinition>();
 
   const decorateMcpToolCandidate = (candidate: unknown): void => {
     if (getCustomOverrideForCandidate(candidate)) {
@@ -1974,7 +1991,8 @@ export function registerToolDisplayOverrides(
     }
 
     const toolName = getTextField(candidate, "name");
-    if (!toolName || wrappedMcpToolNames.has(toolName)) {
+    const runtimeTool = candidate as RuntimeToolDefinition;
+    if (!toolName || decoratedMcpTools.has(runtimeTool)) {
       return;
     }
 
@@ -2003,8 +2021,7 @@ export function registerToolDisplayOverrides(
             ),
           };
 
-    const runtimeTool = candidate as RuntimeToolDefinition;
-    applyToolDisplayDecorationInPlace(
+    if (!applyToolDisplayDecorationInPlace(
       runtimeTool,
       toolDisplayApi,
       {
@@ -2022,7 +2039,9 @@ export function registerToolDisplayOverrides(
           );
         },
       },
-    );
+    )) {
+      return;
+    }
     Object.assign(runtimeTool, {
       label: toolLabel,
       description: toolDescription,
@@ -2031,7 +2050,7 @@ export function registerToolDisplayOverrides(
       prepareArguments: prepareArgumentsDelegate,
     });
 
-    wrappedMcpToolNames.add(toolName);
+    decoratedMcpTools.add(runtimeTool);
   };
 
   const installMcpRegistrationInterceptor = (): void => {
@@ -2047,14 +2066,14 @@ export function registerToolDisplayOverrides(
       this: ExtensionAPI,
       tool: ToolDefinition,
     ): void {
-      originalRegisterTool.call(this, tool);
       try {
         if (!decorateCustomToolOverrideCandidate(tool)) {
           decorateMcpToolCandidate(tool);
         }
       } catch (error) {
         logToolDisplayDebug("Tool display registration decoration failed.", error);
       }
+      originalRegisterTool.call(this, tool);
     } as ExtensionAPI["registerTool"];
 
     pi.registerTool = wrappedRegisterTool;

src/types.ts

@@ -42,6 +42,7 @@ export interface CustomToolOverrideConfig {
 	enabled: boolean;
 	kind: CustomToolOverrideKind;
 	outputMode: CustomToolOutputMode;
+	preserveCallRenderer?: boolean;
 }
 
 export interface ToolDisplayConfig {