@pi-archimedes/ui 2.9.0

package/src/codemode/renderer.ts

@@ -0,0 +1,399 @@
+import { highlightCode } from "@earendil-works/pi-coding-agent";
+import type { Theme } from "@earendil-works/pi-coding-agent";
+import {
+  Container,
+  getImageDimensions,
+  imageFallback,
+  Text,
+  type Component,
+} from "@earendil-works/pi-tui";
+import { renderToolHeader } from "@pi-archimedes/core/tool-render";
+
+/** Nested-call args preview budget, collapsed (the native renderer uses 80). */
+const COLLAPSED_CALL_ARGS_CHARS = 80;
+/** Collapsed call-list budget: the last 5 calls, with an expand hint. */
+const COLLAPSED_CALL_COUNT = 5;
+
+/** The executor's result header, stripped before the output is displayed. */
+const SCRIPT_HEADER = /^Script (completed|failed)\nWall time [\d.]+ seconds\nOutput:\n$/;
+
+export interface CodemodeNestedCall {
+  id: string;
+  name: string;
+  /** Compact JSON of the arguments, truncated for display. */
+  args: string;
+  status: "running" | "ok" | "error" | "cancelled";
+  durationMs?: number;
+  /** Error text, truncated for display. */
+  error?: string;
+  /** Cost in USD of a `models.*` call that reported usage. */
+  cost?: number;
+}
+
+export interface CodemodeToolDetails {
+  calls: CodemodeNestedCall[];
+  /** Temp file with the full text output, when the output was truncated. */
+  fullOutputPath?: string;
+}
+
+export interface CodemodeRendererState {
+  startedAt?: number | undefined;
+  endedAt?: number | undefined;
+  interval?: NodeJS.Timeout | undefined;
+}
+
+const activeStates = new Set<CodemodeRendererState>();
+
+/**
+ * Clear live-timer intervals of in-flight codemode rows. Called on
+ * session_shutdown so a cancelled script does not keep invalidating a
+ * component that no longer exists.
+ */
+export function clearActiveCodemodeIntervals(): void {
+  for (const state of activeStates) {
+    if (state.interval) {
+      clearInterval(state.interval);
+      state.interval = undefined;
+    }
+  }
+  activeStates.clear();
+}
+
+/**
+ * Format a duration in milliseconds: `<X>ms` under a second, `<X.X>s`
+ * otherwise (the native codemode renderer's format).
+ */
+export function formatCodemodeDuration(ms: number | undefined): string {
+  if (ms === undefined) return "";
+  if (ms < 0) ms = 0;
+  return ms < 1000 ? `${Math.round(ms)}ms` : `${(ms / 1000).toFixed(1)}s`;
+}
+
+/**
+ * Cents for larger amounts, two significant digits for the fractions of a
+ * cent classifier calls cost (the native renderer's format).
+ */
+export function formatCost(cost: number): string {
+  return `$${cost >= 0.01 ? cost.toFixed(2) : cost.toPrecision(2)}`;
+}
+
+/**
+ * Status glyph for a nested tool call.
+ */
+function statusIcon(
+  status: CodemodeNestedCall["status"],
+  theme: Theme,
+): string {
+  switch (status) {
+    case "running":
+      return theme.fg("warning", "…");
+    case "ok":
+      return theme.fg("success", "✓");
+    case "error":
+      return theme.fg("error", "✗");
+    case "cancelled":
+      return theme.fg("muted", "⊘");
+  }
+}
+
+/**
+ * One nested tool call row:
+ * `<glyph> <name> <args> <duration> <cost>` — args truncated to
+ * COLLAPSED_CALL_ARGS_CHARS when not expanded; an expanded failed call
+ * appends its error text indented below.
+ */
+export function formatNestedCall(
+  call: CodemodeNestedCall,
+  theme: Theme,
+  expanded: boolean,
+): string {
+  const args =
+    !expanded && call.args.length > COLLAPSED_CALL_ARGS_CHARS
+      ? `${call.args.slice(0, COLLAPSED_CALL_ARGS_CHARS - 3)}...`
+      : call.args;
+  let line = `${statusIcon(call.status, theme)} ${theme.fg("toolTitle", call.name)}`;
+  if (args) line += ` ${theme.fg("muted", args)}`;
+  const duration = formatCodemodeDuration(call.durationMs);
+  if (duration) line += ` ${theme.fg("dim", duration)}`;
+  if (call.cost) line += ` ${theme.fg("dim", formatCost(call.cost))}`;
+  if (expanded && call.error) {
+    line += `\n    ${theme.fg("error", call.error.split("\n").join("\n    "))}`;
+  }
+  return line;
+}
+
+/**
+ * The script output as display text: text blocks joined, the executor's
+ * "Script completed/failed" header dropped, and a fallback indicator line
+ * for image blocks that cannot be shown inline.
+ */
+export function getCodemodeOutput(
+  content: Array<{ type: string; text?: string; data?: string; mimeType?: string }>,
+): string {
+  const textBlocks = content.filter(
+    (b) => b.type === "text" && typeof b.text === "string",
+  );
+  const [first, ...rest] = textBlocks;
+  const hasHeader = first !== undefined && SCRIPT_HEADER.test(first.text!);
+  const text = (hasHeader ? rest : textBlocks).map((b) => b.text!).join("\n");
+  // The renderer is text-only: image blocks are never rendered as images,
+  // so they ALWAYS get a fallback indicator (mime + dimensions when known),
+  // regardless of the terminal's inline-image support.
+  const imageBlocks = content.filter((b) => b.type === "image");
+  if (imageBlocks.length > 0) {
+    const indicators = imageBlocks
+      .map((img) => {
+        const mimeType = img.mimeType ?? "image/unknown";
+        const dims =
+          img.data && img.mimeType
+            ? (getImageDimensions(img.data, img.mimeType) ?? undefined)
+            : undefined;
+        return imageFallback(mimeType, dims);
+      })
+      .join("\n");
+    return text ? `${text}\n${indicators}` : indicators;
+  }
+  return text;
+}
+
+/**
+ * Track the script's wall clock for the collapsed row: start on the first
+ * render after execution began, end when the result settles (or errors).
+ * Mirrors the bash renderer's timer lifecycle.
+ *
+ * `isPartial` means "still in flight": the CALL renderer passes `true`
+ * (a call is in flight until its result settles) — passing `false` there
+ * would record `endedAt` at call-render time and freeze the elapsed timer
+ * for the rest of the script. The interval is (re)created only while the
+ * result has not settled yet (`endedAt === undefined`), so a re-render of a
+ * settled call never spawns a permanent interval.
+ */
+function updateTiming(
+  ctx: {
+    executionStarted?: boolean;
+    isError?: boolean;
+    state?: CodemodeRendererState;
+    invalidate?: () => void;
+  } | undefined,
+  isPartial: boolean,
+): void {
+  const state: CodemodeRendererState = ctx?.state ?? {};
+  if (ctx && !ctx.state) ctx.state = state;
+
+  if ((ctx?.executionStarted || isPartial) && state.startedAt === undefined) {
+    state.startedAt = Date.now();
+  }
+  if (isPartial && !state.interval && state.endedAt === undefined) {
+    if (ctx?.invalidate) {
+      state.interval = setInterval(() => {
+        ctx.invalidate?.();
+      }, 1000);
+      activeStates.add(state);
+    }
+  }
+  if (!isPartial || ctx?.isError) {
+    state.endedAt ??= Date.now();
+    if (state.interval) {
+      clearInterval(state.interval);
+      state.interval = undefined;
+    }
+    activeStates.delete(state);
+  }
+}
+
+/**
+ * The script's wall time in ms, parsed from the result header when present
+ * (more accurate than the renderer's timer, which spans TUI updates).
+ */
+export function parseWallTimeMs(
+  content: Array<{ type: string; text?: string }>,
+): number | undefined {
+  const first = content[0];
+  if (first?.type !== "text" || typeof first.text !== "string") return undefined;
+  const match = first.text.match(
+    /^Script (completed|failed)\nWall time ([\d.]+) seconds/,
+  );
+  if (!match) return undefined;
+  const seconds = Number.parseFloat(match[2]!);
+  return Number.isFinite(seconds) ? seconds * 1000 : undefined;
+}
+
+/**
+ * Render the codemode tool call: the bold `codemode` header. Collapsed, the
+ * header stands alone — the script (syntax-highlighted) shows when expanded.
+ */
+export function renderCodemodeCall(
+  args: unknown,
+  theme: Theme,
+  context: unknown,
+): Component {
+  const ctx = context as {
+    executionStarted?: boolean;
+    state?: CodemodeRendererState;
+    args?: { code?: string };
+  } | undefined;
+
+  // The call is in flight until its result settles — pass `true` so the
+  // timer is not finalized at call-render time (see updateTiming).
+  updateTiming(ctx, true);
+
+  // Native semantics: a missing or empty `code` renders the header alone;
+  // a present-but-non-string `code` renders the invalid-arg marker.
+  const rawCode =
+    (args as { code?: unknown } | undefined)?.code ?? ctx?.args?.code;
+  let code: string | null;
+  if (typeof rawCode === "string") {
+    code = rawCode;
+  } else if (rawCode === undefined || rawCode === null) {
+    code = "";
+  } else {
+    code = null;
+  }
+
+  const container = new Container();
+  const title = renderToolHeader("codemode", undefined, theme);
+
+  if (code === null) {
+    container.addChild(new Text(`${title} ${theme.fg("error", "[invalid arg]")}`, 0, 0));
+    return container;
+  }
+
+  container.addChild(new Text(title, 0, 0));
+
+  // Empty script: header alone (the native renderer's behavior).
+  if (code.length === 0) return container;
+
+  // Collapsed: header alone — the script is the expandable detail.
+  if (!(context as { expanded?: boolean }).expanded) return container;
+
+  // The `// @options:` line is part of the script, so options show as-is.
+  const highlighted = highlightCode(
+    code.replace(/\t/g, "   ").replace(/\r/g, "").trimEnd(),
+    "javascript",
+  ).join("\n");
+  container.addChild(new Text(highlighted, 0, 0));
+  return container;
+}
+
+/**
+ * Render the codemode tool result.
+ *
+ * Collapsed: the nested tool calls (last 5, args truncated) — the script's
+ * activity at a glance, no summary line. A script that made no calls falls
+ * back to one `<glyph> <duration>` line so its outcome is still visible.
+ * Expanded: every nested call (with per-call status, duration, and model
+ * cost), the script output without the executor's header, the full-output
+ * path when truncated, and the timing + status.
+ */
+export function renderCodemodeResult(
+  result: unknown,
+  options: { expanded?: boolean; isPartial?: boolean },
+  theme: Theme,
+  context: unknown,
+): Component {
+  const ctx = context as {
+    executionStarted?: boolean;
+    isError?: boolean;
+    state?: CodemodeRendererState;
+    invalidate?: () => void;
+  } | undefined;
+
+  updateTiming(ctx, Boolean(options.isPartial));
+
+  const res = result as {
+    content?: Array<{ type: string; text?: string; data?: string; mimeType?: string }>;
+    details?: CodemodeToolDetails;
+  } | undefined;
+  const content = Array.isArray(res?.content) ? res.content : [];
+  const calls = res?.details?.calls ?? [];
+  const isError = Boolean(ctx?.isError);
+
+  const state: CodemodeRendererState = ctx?.state ?? {};
+  const now = Date.now();
+  const wallTime = parseWallTimeMs(content);
+  const timerMs =
+    state.startedAt !== undefined
+      ? Math.max(0, (state.endedAt ?? now) - state.startedAt)
+      : undefined;
+  const durationMs = wallTime ?? timerMs;
+
+  const container = new Container();
+
+  // Collapsed view.
+  if (!options.expanded) {
+    if (calls.length > 0) {
+      const shown = calls.slice(-COLLAPSED_CALL_COUNT);
+      const callLines = shown.map((call) => formatNestedCall(call, theme, false));
+      if (shown.length < calls.length) {
+        callLines.unshift(
+          theme.fg(
+            "muted",
+            `... (${calls.length - shown.length} earlier calls, ctrl+o to expand)`,
+          ),
+        );
+      }
+      // A failure must never be hidden behind success glyphs: when the
+      // script failed after successful calls, append the failure marker.
+      if (!options.isPartial && isError) {
+        callLines.push(theme.fg("error", "✗ Script failed"));
+      }
+      container.addChild(new Text(callLines.join("\n"), 0, 0));
+    } else {
+      // No calls: a single status line so the script's outcome is visible.
+      const glyph = options.isPartial
+        ? theme.fg("warning", "▸")
+        : isError
+          ? theme.fg("error", "✗")
+          : theme.fg("success", "✓");
+      const duration = formatCodemodeDuration(durationMs);
+      container.addChild(new Text(`${glyph} ${theme.fg("muted", duration)}`, 0, 0));
+    }
+    return container;
+  }
+
+  // Expanded view.
+  const parts: string[] = [];
+
+  if (calls.length > 0) {
+    const lines = calls.map((call) => formatNestedCall(call, theme, true));
+    const priced = calls.filter((call) => call.cost);
+    if (priced.length > 1) {
+      const total = priced.reduce((sum, call) => sum + (call.cost ?? 0), 0);
+      lines.push(theme.fg("muted", `Model calls: ${formatCost(total)}`));
+    }
+    parts.push(lines.join("\n"));
+  }
+
+  if (!options.isPartial) {
+    const output = getCodemodeOutput(content).trim();
+    if (output.length > 0) {
+      const color = isError ? "error" : "toolOutput";
+      parts.push(
+        output.split("\n").map((line) => theme.fg(color, line)).join("\n"),
+      );
+    }
+    if (res?.details?.fullOutputPath) {
+      parts.push(theme.fg("muted", `Full output: ${res.details.fullOutputPath}`));
+    }
+  }
+
+  if (durationMs !== undefined) {
+    const label = options.isPartial ? "Elapsed" : "Took";
+    parts.push(
+      theme.fg("muted", `${label} ${formatCodemodeDuration(durationMs)}`),
+    );
+  }
+
+  if (!options.isPartial) {
+    parts.push(
+      isError
+        ? theme.fg("error", "✗ Script failed")
+        : theme.fg("success", "✓ Done"),
+    );
+  }
+
+  if (parts.length > 0) {
+    container.addChild(new Text(parts.join("\n\n"), 0, 0));
+  }
+  return container;
+}

package/src/codemode/tool.test.ts

@@ -0,0 +1,149 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { registerCodemodeToolOverride } from "./tool.js";
+
+const nativeSchema = { type: "object", properties: { code: {} } };
+const defFromFactory = {
+  name: "codemode",
+  label: "codemode",
+  description: "native description",
+  parameters: { type: "object", properties: { code: {} } },
+  execute: async () => ({}),
+  prepareLoadout: (loadout: unknown) => loadout,
+  exposure: "model-only",
+  constrainedSampling: { type: "grammar" },
+};
+
+function makePi(
+  overrides: {
+    allTools?: Array<{ name: string; parameters?: unknown; namespace?: unknown }>;
+    getSettings?: () => {
+      codemode?: { mode?: string; inlineBudget?: unknown };
+    };
+  } = {},
+) {
+  const registerTool = vi.fn();
+  const appendEntry = vi.fn();
+  const pi = {
+    registerTool,
+    appendEntry,
+    getAllTools: vi.fn(() => overrides.allTools ?? []),
+    ...(overrides.getSettings ? { getSettings: overrides.getSettings } : {}),
+  };
+  return { pi, registerTool, appendEntry } as const;
+}
+
+type MockModule = {
+  createCodemodeToolDefinition: (options: Record<string, unknown>) => unknown;
+  codemodeSchema: unknown;
+};
+
+// `vi.mock` is hoisted above the imports, so the mutable module handle must
+// live in a hoisted scope (referencing a plain `let` would hit the TDZ when
+// the mock factory first runs during the static import of ./tool.js).
+const state = vi.hoisted(() => ({
+  module: undefined as MockModule | undefined,
+}));
+
+vi.mock("./loader.js", () => ({
+  loadCodemodeModule: vi.fn(async () => state.module),
+}));
+
+beforeEach(() => {
+  state.module = {
+    createCodemodeToolDefinition: (options) => ({ ...defFromFactory, options }),
+    codemodeSchema: nativeSchema,
+  };
+});
+
+describe("registerCodemodeToolOverride", () => {
+  it("registers the native definition with the Archimedes renderers", async () => {
+    const { pi, registerTool } = makePi();
+    const ok = await registerCodemodeToolOverride(pi as never);
+
+    expect(ok).toBe(true);
+    expect(registerTool).toHaveBeenCalledTimes(1);
+    const registered = registerTool.mock.calls[0]![0] as Record<string, unknown>;
+    // Native definition fields are preserved.
+    expect(registered.name).toBe("codemode");
+    expect(registered.execute).toBe(defFromFactory.execute);
+    expect(registered.prepareLoadout).toBe(defFromFactory.prepareLoadout);
+    expect(registered.exposure).toBe("model-only");
+    // The built-in registers codemode inactive; the override must match.
+    expect(registered.defaultActive).toBe(false);
+    // The presentation is replaced.
+    expect(typeof registered.renderCall).toBe("function");
+    expect(typeof registered.renderResult).toBe("function");
+  });
+
+  it("preserves the running CLI's schema reference when the built-in is still registered", async () => {
+    const { pi, registerTool } = makePi({
+      allTools: [{ name: "codemode", parameters: nativeSchema }],
+    });
+    await registerCodemodeToolOverride(pi as never);
+
+    const registered = registerTool.mock.calls[0]![0] as Record<string, unknown>;
+    expect(registered.parameters).toBe(nativeSchema);
+  });
+
+  it("keeps the factory schema when no built-in registration is visible", async () => {
+    const { pi, registerTool } = makePi();
+    await registerCodemodeToolOverride(pi as never);
+
+    const registered = registerTool.mock.calls[0]![0] as Record<string, unknown>;
+    expect(registered.parameters).toBe(defFromFactory.parameters);
+  });
+
+  it("wires appendEntry, models, and the settings-backed options", async () => {
+    const settings = { codemode: { mode: "only", inlineBudget: 1500 } };
+    const { pi, registerTool, appendEntry } = makePi({
+      getSettings: () => settings,
+    });
+    await registerCodemodeToolOverride(pi as never);
+
+    const registered = registerTool.mock.calls[0]![0] as Record<string, unknown>;
+    const options = (registered as { options?: Record<string, unknown> }).options!;
+    const getMode = options.getMode as () => string;
+    const getInlineBudget = options.getInlineBudget as () => number | undefined;
+    expect(options.models).toBe(true);
+    expect(typeof options.appendEntry).toBe("function");
+    expect(typeof options.getToolNamespace).toBe("function");
+    expect(getMode()).toBe("only");
+    expect(getInlineBudget()).toBe(1500);
+
+    // appendEntry forwards to pi.appendEntry.
+    await (options.appendEntry as (t: string, d: unknown) => void)(
+      "codemode-store",
+      { set: {} },
+    );
+    expect(appendEntry).toHaveBeenCalledWith("codemode-store", { set: {} });
+  });
+
+  it("defaults mode to on and inlineBudget to undefined without settings", async () => {
+    const { pi, registerTool } = makePi();
+    await registerCodemodeToolOverride(pi as never);
+
+    const registered = registerTool.mock.calls[0]![0] as Record<string, unknown>;
+    const options = (registered as { options?: Record<string, unknown> }).options!;
+    expect((options.getMode as () => string)()).toBe("on");
+    expect((options.getInlineBudget as () => number | undefined)()).toBeUndefined();
+  });
+
+  it("treats a non-finite inlineBudget as undefined", async () => {
+    const { pi, registerTool } = makePi({
+      getSettings: () => ({ codemode: { inlineBudget: Number.NaN } }),
+    });
+    await registerCodemodeToolOverride(pi as never);
+
+    const registered = registerTool.mock.calls[0]![0] as Record<string, unknown>;
+    const options = (registered as { options?: Record<string, unknown> }).options!;
+    expect((options.getInlineBudget as () => number | undefined)()).toBeUndefined();
+  });
+
+  it("returns false and registers nothing when the module cannot be loaded", async () => {
+    state.module = undefined;
+    const { pi, registerTool } = makePi();
+    const ok = await registerCodemodeToolOverride(pi as never);
+    expect(ok).toBe(false);
+    expect(registerTool).not.toHaveBeenCalled();
+  });
+});

package/src/codemode/tool.ts

@@ -0,0 +1,96 @@
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { loadCodemodeModule } from "./loader.js";
+import {
+  renderCodemodeCall,
+  renderCodemodeResult,
+  clearActiveCodemodeIntervals,
+} from "./renderer.js";
+
+export { clearActiveCodemodeIntervals } from "./renderer.js";
+
+/**
+ * Register an Archimedes-styled override of the `codemode` tool.
+ *
+ * `codemode` is a replaceable pi built-in extension: registering a tool with
+ * the same name takes over completely (pi omits the built-in one), so the
+ * override must reproduce the built-in's full definition — the QuickJS
+ * sandbox executor, the dynamic `prepareLoadout` description, the
+ * `model-only` exposure, and `defaultActive: false` (the built-in registers
+ * codemode inactive; it is activated via `--tools`, `defaultTools`, or MCP
+ * auto-enable). Only the presentation is replaced.
+ *
+ * The definition is built from the running CLI's own copy of the codemode
+ * module (see {@link loadCodemodeModule}) so the executor matches what the
+ * CLI expects. The `parameters` reference is preserved from the built-in's
+ * registration when it is still visible (`pi.getAllTools()`): the CLI's
+ * MCP and tool-search extensions recognise the codemode tool by schema
+ * identity (`isCodemodeTool`), and a fresh schema object from a second
+ * module instance would defeat those checks.
+ *
+ * Returns `false` when the codemode module cannot be loaded (e.g. a pi
+ * version without the codemode extension) — pi's native rendering then
+ * stands in.
+ *
+ * Note: the take-over makes pi print a startup notice ("built-in extension
+ * `codemode` was not loaded") — the same notice a third-party MCP extension
+ * triggers when it replaces the built-in MCP extension. It is informational:
+ * the tool still works, now with this package's renderers. Set
+ * `codemodeToolStyling` to `false` to run the native tool instead.
+ */
+export async function registerCodemodeToolOverride(
+  pi: ExtensionAPI,
+): Promise<boolean> {
+  const mod = await loadCodemodeModule();
+  if (!mod) return false;
+
+  const def = mod.createCodemodeToolDefinition({
+    appendEntry: (customType: string, data: unknown) =>
+      pi.appendEntry(customType, data),
+    models: true,
+    getToolNamespace: (name: string) =>
+      (pi.getAllTools().find((t) => t.name === name) as {
+        namespace?: unknown;
+      } | undefined)?.namespace,
+    // `getSettings` is newer than the oldest pi this package type-checks
+    // against; read it defensively (the built-in extension's wiring).
+    getMode: () =>
+      readSettings(pi)?.codemode?.mode === "only" ? "only" : "on",
+    getInlineBudget: () => {
+      const budget = readSettings(pi)?.codemode?.inlineBudget;
+      return typeof budget === "number" && Number.isFinite(budget) && budget >= 0
+        ? budget
+        : undefined;
+    },
+  }) as Record<string, unknown>;
+
+  // The built-in (when loaded) is still registered at this point: keep its
+  // schema object reference so `isCodemodeTool` identity checks keep working.
+  const native = pi.getAllTools().find((t) => t.name === "codemode");
+  if (native?.parameters !== undefined) {
+    def.parameters = native.parameters;
+  }
+
+  pi.registerTool({
+    ...def,
+    renderCall: renderCodemodeCall,
+    renderResult: renderCodemodeResult,
+    defaultActive: false,
+  } as never);
+  return true;
+}
+
+function readSettings(
+  pi: ExtensionAPI,
+): { codemode?: { mode?: string; inlineBudget?: unknown } } | undefined {
+  const getSettings = (
+    pi as unknown as { getSettings?: () => unknown }
+  ).getSettings;
+  if (typeof getSettings !== "function") return undefined;
+  try {
+    return getSettings.call(pi) as
+      | { codemode?: { mode?: string; inlineBudget?: unknown } }
+      | undefined;
+  } catch {
+    return undefined;
+  }
+}