@pi-kaush/pi-tool-call-markers 0.2.7 → 0.2.8

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@pi-kaush/pi-tool-call-markers",
-  "version": "0.2.7",
+  "version": "0.2.8",
   "description": "Compact and group Pi tool calls, with a bundled extension for adjacent thinking blocks.",
   "license": "MIT",
   "author": "Kaushik Gopal",

package/src/index.ts

@@ -61,6 +61,7 @@ type PresentationPatchState = {
   theme?: ThemeLike;
   collapseParallel: boolean;
   groupCache: WeakMap<ToolExecutionRow, GroupRenderCache>;
+  collapsedCache: WeakMap<ToolExecutionRow, CollapsedRenderCache>;
   rowVersions: WeakMap<ToolExecutionRow, number>;
   rowGroups: WeakMap<ToolExecutionRow, ToolExecutionRow[]>;
   rowSignatures: WeakMap<ToolExecutionRow, string>;
@@ -85,6 +86,14 @@ type GroupRenderCache = {
   width: number;
 };
 
+type CollapsedRenderCache = {
+  lines: string[];
+  signature: string;
+  version: number;
+  themeSample: string;
+  width: number;
+};
+
 function envEnabled(name: string, defaultValue: boolean): boolean {
   const value = process.env[name]?.trim().toLowerCase();
   if (!value) return defaultValue;
@@ -116,6 +125,33 @@ function hasVisibleContent(line: string): boolean {
   return stripAnsi(line).trim().length > 0;
 }
 
+// State-transition fingerprint used by both the updateDisplay version bump
+// and the settled-row render cache. It deliberately covers only shape
+// transitions (partial/expanded/result/error): while a row is partial its
+// content streams and it is never cached, and once settled the rendered
+// content is fully determined by this fingerprint + width + theme.
+function rowSignatureOf(row: ToolExecutionRow): string {
+  return `${row.isPartial}|${row.expanded}|${row.result ? 1 : 0}|${row.result?.isError ? 1 : 0}`;
+}
+
+// Theme samples are five theme.fg calls; computing them per row per frame
+// would reintroduce the per-keystroke cost the render caches remove, so
+// cache by theme object identity (a theme switch swaps the object).
+const themeSamples = new WeakMap<ThemeLike, string>();
+function themeSampleFor(theme: ThemeLike): string {
+  let sample = themeSamples.get(theme);
+  if (!sample) {
+    sample =
+      theme.fg("toolTitle", "x") +
+      theme.fg("muted", "x") +
+      theme.fg("dim", "x") +
+      theme.fg("warning", "x") +
+      theme.fg("error", "x");
+    themeSamples.set(theme, sample);
+  }
+  return sample;
+}
+
 type InsetLayout = {
   contentWidth: number;
   left: number;
@@ -1115,12 +1151,7 @@ function renderGroupedToolRows(
 ): string[] {
   const theme = state.theme;
   if (!theme) return state.originalRender.call(row, width);
-  const themeSample =
-    theme.fg("toolTitle", "x") +
-    theme.fg("muted", "x") +
-    theme.fg("dim", "x") +
-    theme.fg("warning", "x") +
-    theme.fg("error", "x");
+  const themeSample = themeSampleFor(theme);
   const cached = state.groupCache.get(row);
   const hasLiveMembers = rows.some(isLiveRow);
   if (
@@ -1330,6 +1361,7 @@ function installPresentationPatch(): PresentationPatchState | undefined {
     const state: PresentationPatchState = {
       collapseParallel: envEnabled(COLLAPSE_PARALLEL_ENV, true),
       groupCache: new WeakMap(),
+      collapsedCache: new WeakMap(),
       rowVersions: new WeakMap(),
       rowGroups: new WeakMap(),
       rowSignatures: new WeakMap(),
@@ -1342,7 +1374,7 @@ function installPresentationPatch(): PresentationPatchState | undefined {
       // Group members always bump so their leader's cache refreshes; other
       // rows only bump on state transitions, so bash's per-second invalidate
       // ticks and resize invalidations stop busting caches.
-      const signature = `${this.isPartial}|${this.expanded}|${this.result ? 1 : 0}|${this.result?.isError ? 1 : 0}`;
+      const signature = rowSignatureOf(this);
       if (
         state.rowGroups.has(this) ||
         state.rowSignatures.get(this) !== signature
@@ -1357,18 +1389,72 @@ function installPresentationPatch(): PresentationPatchState | undefined {
       this: ToolExecutionRow,
       width: number,
     ): string[] {
-      const lines = state.originalRender.call(this, width);
+      const theme = state.theme;
       if (
         typeof this.expanded !== "boolean" ||
         typeof this.isPartial !== "boolean" ||
         this.expanded ||
-        !state.theme ||
-        (lines.length === 0 && (this.imageComponents?.length ?? 0) === 0)
+        !theme
       ) {
+        return state.originalRender.call(this, width);
+      }
+
+      // ================================================================
+      // Settled-row render cache — PLEASE DO NOT REMOVE THIS FAST PATH.
+      // ================================================================
+      // Pi's TUI re-renders the whole transcript on every keystroke, and
+      // the collapsed-row decoration below rebuilds strings for every row
+      // it touches. On a long session (~1,000 messages, hundreds of tool
+      // rows) that measured at ~80 ms per keystroke — visible input lag.
+      //
+      // Settled rows (not partial, not expanded) have fully static render
+      // output, so their collapsed lines are cached keyed by width +
+      // rowSignatureOf + rowVersions + theme sample. updateDisplay bumps
+      // rowVersions on every real state transition, and partial rows
+      // bypass the cache entirely because their content streams. This
+      // mirrors the existing groupCache discipline; keep them in sync.
+      // ================================================================
+      if (!this.isPartial) {
+        const signature = rowSignatureOf(this);
+        const version = state.rowVersions.get(this) ?? 0;
+        const themeSample = themeSampleFor(theme);
+        const cached = state.collapsedCache.get(this);
+        if (
+          cached &&
+          cached.width === width &&
+          cached.signature === signature &&
+          cached.version === version &&
+          cached.themeSample === themeSample
+        ) {
+          return cached.lines;
+        }
+        const lines = state.originalRender.call(this, width);
+        if (lines.length === 0 && (this.imageComponents?.length ?? 0) === 0) {
+          return lines;
+        }
+        try {
+          const collapsed = renderCollapsedToolRow(this, width, theme);
+          state.collapsedCache.set(this, {
+            lines: collapsed,
+            signature,
+            version,
+            themeSample,
+            width,
+          });
+          return collapsed;
+        } catch {
+          return lines;
+        }
+      }
+
+      // Partial (streaming) rows keep the collapsed presentation but are
+      // never cached: their content changes with every streamed chunk.
+      const lines = state.originalRender.call(this, width);
+      if (lines.length === 0 && (this.imageComponents?.length ?? 0) === 0) {
         return lines;
       }
       try {
-        return renderCollapsedToolRow(this, width, state.theme);
+        return renderCollapsedToolRow(this, width, theme);
       } catch {
         return lines;
       }

package/src/thinking-block-merger.ts

@@ -5,6 +5,176 @@ const THINKING_GROUPING_PATCHED = Symbol.for("kg.pi.thinkingGrouping.v1");
 const PI_SPINNER_FRAMES = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
 const SPINNER_INTERVAL_MS = 80;
 
+// Visual experiment for the settled "+ Thought" label. "inherit" keeps Pi's
+// native styling (italic + thinkingText); "gray" sits halfway between the
+// muted and text theme colors. The env var overrides the default so the
+// variants can be compared without editing code or republishing.
+const THOUGHT_LABEL_COLOR_ENV = "PI_TOOL_CALL_MARKERS_THOUGHT_COLOR";
+type ThoughtLabelColor = "inherit" | "orange" | "gray";
+const DEFAULT_THOUGHT_LABEL_COLOR: ThoughtLabelColor = "orange";
+// cobalt2's orange token (#ffb86c).
+const THOUGHT_LABEL_ORANGE = { r: 255, g: 184, b: 108 };
+
+function thoughtLabelColorChoice(): ThoughtLabelColor {
+  const raw = process.env[THOUGHT_LABEL_COLOR_ENV];
+  return raw === "inherit" || raw === "orange" || raw === "gray"
+    ? raw
+    : DEFAULT_THOUGHT_LABEL_COLOR;
+}
+
+const CUBE_VALUES = [0, 95, 135, 175, 215, 255];
+const BASIC_ANSI_RGB: Array<[number, string] | undefined> = [
+  [0, "#000000"],
+  [1, "#800000"],
+  [2, "#008000"],
+  [3, "#808000"],
+  [4, "#000080"],
+  [5, "#800080"],
+  [6, "#008080"],
+  [7, "#c0c0c0"],
+  [8, "#808080"],
+  [9, "#ff0000"],
+  [10, "#00ff00"],
+  [11, "#ffff00"],
+  [12, "#0000ff"],
+  [13, "#ff00ff"],
+  [14, "#00ffff"],
+  [15, "#ffffff"],
+];
+
+type Rgb = { r: number; g: number; b: number };
+
+type ThemeDetail = {
+  fg(color: string, text: string): string;
+  getFgAnsi?(color: string): string;
+  getColorMode?(): "truecolor" | "256color" | string;
+};
+
+let activeTheme: ThemeDetail | undefined;
+
+function hexToRgb(hex: string): Rgb | undefined {
+  const match = /^#([0-9a-f]{6})$/i.exec(hex.trim());
+  if (!match) return undefined;
+  const value = parseInt(match[1]!, 16);
+  return { r: value >> 16, g: (value >> 8) & 255, b: value & 255 };
+}
+
+// Mirrors Pi's own mapping so a named color can round-trip through a
+// 256-color terminal (e.g. when COLORFGBG forces 256color mode).
+function ansi256ToRgb(index: number): Rgb | undefined {
+  if (index < 16) {
+    const hex = BASIC_ANSI_RGB[index]?.[1];
+    return hex ? hexToRgb(hex) : undefined;
+  }
+  if (index < 232) {
+    const cube = index - 16;
+    const channel = (n: number) => (n === 0 ? 0 : 55 + n * 40);
+    return {
+      r: channel(Math.floor(cube / 36)),
+      g: channel(Math.floor((cube % 36) / 6)),
+      b: channel(cube % 6),
+    };
+  }
+  const gray = 8 + (index - 232) * 10;
+  return { r: gray, g: gray, b: gray };
+}
+
+function rgbToAnsi256({ r, g, b }: Rgb): number {
+  const nearest = (value: number) =>
+    CUBE_VALUES.reduce(
+      (best, candidate) =>
+        Math.abs(candidate - value) < Math.abs(best - value) ? candidate : best,
+      0,
+    );
+  const cubeR = nearest(r);
+  const cubeG = nearest(g);
+  const cubeB = nearest(b);
+  const cubeDistance =
+    (r - cubeR) ** 2 * 0.299 +
+    (g - cubeG) ** 2 * 0.587 +
+    (b - cubeB) ** 2 * 0.114;
+  const luminance = Math.round(0.299 * r + 0.587 * g + 0.114 * b);
+  const grayStep = Math.max(0, Math.min(23, Math.round((luminance - 8) / 10)));
+  const grayValue = 8 + grayStep * 10;
+  const grayDistance = (luminance - grayValue) ** 2;
+  const spread = Math.max(r, g, b) - Math.min(r, g, b);
+  if (spread < 10 && grayDistance < cubeDistance) return 232 + grayStep;
+  return (
+    16 +
+    36 * CUBE_VALUES.indexOf(cubeR) +
+    6 * CUBE_VALUES.indexOf(cubeG) +
+    CUBE_VALUES.indexOf(cubeB)
+  );
+}
+
+function parseAnsiFgRgb(ansi: string): Rgb | undefined {
+  if (/^\x1b\[39m$/.test(ansi)) return undefined;
+  const truecolor = /^\x1b\[38;2;(\d+);(\d+);(\d+)m$/.exec(ansi);
+  if (truecolor) {
+    return {
+      r: Number(truecolor[1]),
+      g: Number(truecolor[2]),
+      b: Number(truecolor[3]),
+    };
+  }
+  const indexed = /^\x1b\[38;5;(\d+)m$/.exec(ansi);
+  if (indexed) return ansi256ToRgb(Number(indexed[1]));
+  const basic = /^\x1b\[(9[0-7]|3[0-7])m$/.exec(ansi);
+  if (basic) {
+    const code = Number(basic[1]);
+    return ansi256ToRgb(code >= 90 ? code - 90 + 8 : code - 30);
+  }
+  return undefined;
+}
+
+function fgRgbAnsi({ r, g, b }: Rgb): string {
+  if (activeTheme?.getColorMode?.() === "256color") {
+    return `\x1b[38;5;${rgbToAnsi256({ r, g, b })}m`;
+  }
+  return `\x1b[38;2;${r};${g};${b}m`;
+}
+
+function midpointRgb(a: Rgb, b: Rgb): Rgb {
+  return {
+    r: Math.round((a.r + b.r) / 2),
+    g: Math.round((a.g + b.g) / 2),
+    b: Math.round((a.b + b.b) / 2),
+  };
+}
+
+let thoughtLabelPrefix = "";
+let thoughtLabelSuffix = "";
+
+function updateThoughtLabelStyle(): void {
+  thoughtLabelPrefix = "";
+  thoughtLabelSuffix = "";
+  const choice = thoughtLabelColorChoice();
+  if (choice === "inherit" || !activeTheme) return;
+
+  let rgb: Rgb | undefined;
+  if (choice === "orange") {
+    rgb = THOUGHT_LABEL_ORANGE;
+  } else if (activeTheme.getFgAnsi) {
+    const muted = parseAnsiFgRgb(activeTheme.getFgAnsi("muted"));
+    const text = parseAnsiFgRgb(activeTheme.getFgAnsi("text"));
+    if (muted && text) rgb = midpointRgb(muted, text);
+  }
+  if (!rgb) return;
+
+  // The styled label replaces Pi's italicized Text node wholesale (see
+  // restyleHiddenThinkingLabel). The leading italic-off still matters: the
+  // TUI's diff renderer can skip bytes shared with the previously drawn
+  // italic line, leaving the terminal in italic state otherwise.
+  thoughtLabelPrefix = `\x1b[23m${fgRgbAnsi(rgb)}`;
+  thoughtLabelSuffix = "\x1b[39m";
+}
+
+export function visibleThoughtLabel(label: string): string {
+  if (!thoughtLabelPrefix) return label;
+  const raw = label.replace(/\x1b\[[0-9;]*m/g, "");
+  return `${thoughtLabelPrefix}${raw}${thoughtLabelSuffix}`;
+}
+
 type AssistantMessageLike = {
   content?: unknown[];
 };
@@ -12,9 +182,46 @@ type AssistantMessageLike = {
 type AssistantMessageRow = {
   hiddenThinkingLabel?: unknown;
   hideThinkingBlock?: unknown;
+  contentContainer?: { children?: unknown[] };
   updateContent(message: AssistantMessageLike, ...args: unknown[]): void;
 };
 
+type TextLikeChild = {
+  text?: unknown;
+  setText?(text: string): void;
+};
+
+// Pi renders the hidden-thinking label as an italic Text node built from the
+// plain label field. Embedding style codes in the field itself is not
+// enough: the TUI diff renderer reuses the byte prefix shared with the
+// previous (italic) frame, so an in-line italic reset may never reach the
+// terminal. Swap the node's text for a self-contained styled version after
+// each native render pass instead.
+function restyleHiddenThinkingLabel(row: AssistantMessageRow): void {
+  if (!thoughtLabelPrefix) return;
+  if (
+    row.hideThinkingBlock !== true ||
+    typeof row.hiddenThinkingLabel !== "string"
+  ) {
+    return;
+  }
+  const label = row.hiddenThinkingLabel;
+  const children = row.contentContainer?.children;
+  if (!Array.isArray(children)) return;
+  const styled = visibleThoughtLabel(label);
+  for (const child of children) {
+    const textChild = child as TextLikeChild | undefined;
+    if (
+      typeof textChild?.text !== "string" ||
+      typeof textChild.setText !== "function" ||
+      !textChild.text.includes(label)
+    ) {
+      continue;
+    }
+    textChild.setText(styled);
+  }
+}
+
 type ThinkingTiming = {
   finishedAt?: number;
   startedAt: number;
@@ -186,6 +393,11 @@ function installThinkingGroupingPatch():
         // Preserve Pi's native label if its private row shape changes.
       }
       Reflect.apply(state.originalUpdateContent, this, [combined, ...args]);
+      try {
+        restyleHiddenThinkingLabel(this);
+      } catch {
+        // Keep Pi's native label styling if the row shape changes.
+      }
     };
 
     state.patchedUpdateContent = patchedUpdateContent;
@@ -224,5 +436,15 @@ function uninstallThinkingGroupingPatch(
 
 export default function (pi: ExtensionAPI) {
   const patch = installThinkingGroupingPatch();
-  pi.on("session_shutdown", () => uninstallThinkingGroupingPatch(patch));
+  pi.on("session_start", (_event, ctx) => {
+    if (ctx.mode !== "tui") return;
+    activeTheme = ctx.ui.theme as unknown as ThemeDetail;
+    updateThoughtLabelStyle();
+  });
+  pi.on("session_shutdown", () => {
+    activeTheme = undefined;
+    thoughtLabelPrefix = "";
+    thoughtLabelSuffix = "";
+    uninstallThinkingGroupingPatch(patch);
+  });
 }