pi-workflow-engine 0.13.0 → 1.1.0

package/.pi/extensions/pi-workflow-engine/src/ui/workflow-widget.ts

@@ -1,127 +1,345 @@
 import type { Theme } from "@earendil-works/pi-coding-agent";
+import { truncateToWidth, visibleWidth, type Component } from "@earendil-works/pi-tui";
 import type { AgentRowSnapshot, PhaseSnapshot, WorkflowProgressSnapshot } from "../progress-types.ts";
-import { formatWorkflowUsageLine } from "../usage.ts";
-import { agentDetailParts, agentLabelColor, formatCount, formatDuration, statusIcon } from "./workflow-format.ts";
+import { formatCount } from "../text.ts";
+import { compactUsage, dot, fit, GLYPH, joinParts, keyHints, progressBar, spread, type ThemeColor } from "./kit.ts";
+import { countAgents, formatDuration, formatElapsed, type WorkflowStatusCounts } from "./workflow-format.ts";
 
+/** pi caps string widgets at this many rows; the component keeps to the same budget. */
 const MAX_WIDGET_LINES = 10;
+/** Width used for surfaces that only accept pre-rendered strings (RPC). */
+export const STRING_WIDGET_WIDTH = 100;
 
-export function renderWorkflowWidgetLines(
-  snapshot: WorkflowProgressSnapshot,
-  theme: Theme,
-): string[] {
+const INSPECTOR_COMMAND = "/workflow:inspector";
+const MAX_BAR_WIDTH = 24;
+const MIN_BAR_WIDTH = 8;
+const MAX_PHASE_COLUMN = 18;
+const MAX_LABEL_COLUMN = 26;
+const FAILURE_LOG = /\b(fail(ed|ure|ing)?|error)\b/i;
+
+let inspectorShortcut: string | null = null;
+
+/** The widget footer names the configured inspector shortcut; registration sets it once from the shortcut config. */
+export function setWorkflowInspectorShortcut(shortcut: string | null): void {
+  inspectorShortcut = shortcut;
+}
+
+function workflowInspectorHint(): string {
+  return inspectorShortcut ?? INSPECTOR_COMMAND;
+}
+
+/**
+ * The live run widget, every line fitted to `width`:
+ * header (name, bar, finished/total, elapsed, usage) → phases in chronological order, the active
+ * phase expanded into agent rows → upcoming declared phases → counters with the inspector hint.
+ */
+export function renderWorkflowWidget(snapshot: WorkflowProgressSnapshot, width: number, theme: Theme): string[] {
+  const now = Date.now();
+  const inner = Math.max(1, width - 2);
   const counts = countAgents(snapshot.phases);
-  const active = counts.running + counts.queued;
-  const elapsed = formatDuration((snapshot.doneAt ?? Date.now()) - snapshot.startedAt);
-  const activity = active > 0 ? theme.fg("accent", "●") : theme.fg("success", "✓");
-
-  const headingParts = [`${counts.done}/${counts.total} done`, elapsed];
-  if (counts.running > 0) headingParts.unshift(`${counts.running} running`);
-  if (counts.queued > 0) headingParts.unshift(`${counts.queued} queued`);
-  if (counts.failed > 0) headingParts.unshift(theme.fg("error", `${counts.failed} failed`));
-
-  const lines: string[] = [
-    `${activity} ${theme.bold(snapshot.title)} ${theme.fg("dim", "·")} ${theme.fg("muted", snapshot.currentPhase)} ${theme.fg("dim", "·")} ${headingParts.join(` ${theme.fg("dim", "·")} `)}`,
-  ];
-
-  const footer = footerLine(snapshot, theme);
-  const bodyBudget = Math.max(0, MAX_WIDGET_LINES - lines.length - (footer ? 1 : 0));
-  const body = visibleBodyLines(snapshot.phases, bodyBudget, theme);
-  const reserveHiddenLine = body.hidden > 0 && bodyBudget > 0;
-  const visibleBodyLinesToRender = reserveHiddenLine ? body.lines.slice(0, Math.max(0, bodyBudget - 1)) : body.lines;
-  const hidden = body.hidden + (body.lines.length - visibleBodyLinesToRender.length);
-
-  lines.push(...visibleBodyLinesToRender);
-  if (hidden > 0 && lines.length < MAX_WIDGET_LINES) {
-    lines.push(`${theme.fg("dim", "└─")} ${theme.fg("dim", `+${hidden} more`)}`);
+  const phases = classifyPhases(snapshot);
+  const failureLog = failureLogLine(snapshot, theme);
+  const footer = footerLine(snapshot, workflowInspectorHint(), inner, theme);
+  const upcoming = upcomingLine(snapshot, theme);
+
+  const bodyBudget = MAX_WIDGET_LINES - 2 - (failureLog ? 1 : 0);
+  const body = phaseBody(phases, upcoming, bodyBudget, inner, now, theme);
+  const lines = [headerLine(snapshot, counts, inner, now, theme), ...body];
+  if (failureLog) lines.push(failureLog);
+  lines.push(footer);
+  return lines.map((line) => ` ${truncateToWidth(line, inner, "…")}`);
+}
+
+/** A TUI widget that re-renders its lines at the terminal's current width on every frame. */
+export class WidthAwareWidget implements Component {
+  constructor(private readonly lines: (width: number) => string[]) {}
+
+  render(width: number): string[] {
+    return this.lines(width);
   }
-  if (footer && lines.length < MAX_WIDGET_LINES) lines.push(footer);
 
-  return lines.slice(0, MAX_WIDGET_LINES);
+  invalidate(): void {}
+}
+
+/** `◆ background · code-review 56b57550 · Find 6/10 · 12s`: one line for a run the user sent to the background. */
+export function renderBackgroundWorkflowLine(snapshot: WorkflowProgressSnapshot, width: number, theme: Theme): string[] {
+  const counts = countAgents(snapshot.phases);
+  const progress = counts.total > 0 ? `${snapshot.currentPhase} ${counts.done + counts.failed}/${counts.total}` : snapshot.currentPhase;
+  const line = joinParts([
+    `${theme.fg("accent", GLYPH.workflow)} ${theme.fg("muted", "background")}`,
+    `${theme.bold(snapshot.title)} ${theme.fg("dim", snapshot.runId.slice(0, 8))}`,
+    theme.fg("muted", progress),
+    counts.failed > 0 ? theme.fg("error", `${counts.failed} failed`) : undefined,
+    theme.fg("dim", formatElapsed(Date.now() - snapshot.startedAt)),
+  ], theme);
+  return [` ${spread(line, keyHints([[workflowInspectorHint(), "inspect"]], theme), Math.max(1, width - 2))}`];
+}
+
+interface ClassifiedPhases {
+  /** Every shown phase in chronological order; finished and active phases interleave when pipelines overlap. */
+  readonly ordered: readonly PhaseSnapshot[];
+  /** The current phase plus any phase that still has running or queued agents. */
+  readonly active: readonly PhaseSnapshot[];
+}
+
+function classifyPhases(snapshot: WorkflowProgressSnapshot): ClassifiedPhases {
+  // The implicit "Workflow" bucket only matters when agents ran before any phase() call.
+  const ordered = snapshot.phases.filter((phase) => phase.agents.length > 0 || phase.title !== "Workflow");
+  const active = ordered.filter((phase) =>
+    phase.title === snapshot.currentPhase || phase.agents.some((agent) => agent.status === "running" || agent.status === "queued"));
+  return { ordered, active };
+}
+
+function headerLine(snapshot: WorkflowProgressSnapshot, counts: WorkflowStatusCounts, width: number, now: number, theme: Theme): string {
+  const name = `${theme.fg("accent", GLYPH.workflow)} ${theme.bold(snapshot.title)}`;
+  const finished = counts.done + counts.failed;
+  const stats = joinParts([
+    counts.total > 0 ? theme.fg("text", `${finished}/${counts.total}`) : undefined,
+    theme.fg("muted", formatElapsed(now - snapshot.startedAt)),
+    compactUsage(snapshot.usage, theme),
+  ], theme);
+  const barWidth = Math.min(MAX_BAR_WIDTH, width - visibleWidth(name) - visibleWidth(stats) - 4);
+  if (counts.total === 0 || barWidth < MIN_BAR_WIDTH) return `${name}  ${stats}`;
+  return `${name}  ${progressBar(finished, counts.total, barWidth, theme)}  ${stats}`;
 }
 
-interface AgentCounts {
-  queued: number;
-  running: number;
-  done: number;
-  failed: number;
-  total: number;
+interface AgentRowPlan {
+  readonly phase: PhaseSnapshot;
+  readonly shown: readonly AgentRowSnapshot[];
+  readonly hidden: number;
 }
 
-function countAgents(phases: readonly PhaseSnapshot[]): AgentCounts {
-  const counts: AgentCounts = { queued: 0, running: 0, done: 0, failed: 0, total: 0 };
-  for (const phase of phases) {
-    for (const agent of phase.agents) {
-      counts[agent.status]++;
-      counts.total++;
+function phaseBody(
+  phases: ClassifiedPhases,
+  upcoming: string | undefined,
+  budget: number,
+  width: number,
+  now: number,
+  theme: Theme,
+): string[] {
+  let ordered: readonly (PhaseSnapshot | "collapsed")[] = phases.ordered;
+  const candidates = phases.active.reduce((total, phase) => total + phase.agents.length, 0);
+  const wantedRows = Math.min(candidates, 3);
+  const fixedLines = () => ordered.length + (upcoming ? 1 : 0);
+  // Keep agent rows visible on long workflows: fold everything finished before the first active phase into one line.
+  if (budget - fixedLines() < wantedRows) {
+    const firstActive = phases.ordered.findIndex((phase) => phases.active.includes(phase));
+    const foldUntil = firstActive === -1 ? phases.ordered.length : firstActive;
+    if (foldUntil > 1) ordered = ["collapsed", ...phases.ordered.slice(foldUntil)];
+  }
+  const showUpcoming = upcoming !== undefined && budget - ordered.length >= 1;
+  const rowBudget = Math.max(0, budget - ordered.length - (showUpcoming ? 1 : 0));
+  const plans = planAgentRows(phases.active, rowBudget);
+
+  const visiblePhases = ordered.filter((entry): entry is PhaseSnapshot => entry !== "collapsed");
+  const titleWidth = Math.min(MAX_PHASE_COLUMN, Math.max(0, ...visiblePhases.map((phase) => visibleWidth(phase.title))));
+  const labelWidth = agentLabelWidth(plans, width);
+  const folded = phases.ordered.slice(0, phases.ordered.length - visiblePhases.length);
+
+  const lines: string[] = [];
+  for (const entry of ordered) {
+    if (entry === "collapsed") {
+      lines.push(foldedPhasesLine(folded, theme));
+      continue;
     }
+    const plan = plans.get(entry);
+    if (!plan) {
+      lines.push(finishedPhaseLine(entry, titleWidth, theme));
+      continue;
+    }
+    lines.push(activePhaseLine(plan, titleWidth, theme));
+    for (const agent of plan.shown) lines.push(agentLine(agent, entry, labelWidth, width, now, theme));
   }
-  return counts;
+  if (showUpcoming && upcoming) lines.push(upcoming);
+  return lines.slice(0, budget);
 }
 
-function visibleBodyLines(phases: readonly PhaseSnapshot[], budget: number, theme: Theme): { lines: string[]; hidden: number } {
-  const lines: string[] = [];
-  let totalRows = 0;
-  const visit = (phase: PhaseSnapshot): void => {
-    if (phase.agents.length === 0 && phase.title === "Workflow") return;
-    const counts = countPhaseAgents(phase);
-    totalRows++;
-    if (lines.length < budget) lines.push(phaseLine(phase.title, counts, theme));
-    appendAgentGroup(phase.agents, (agent) => agent.status === "running" || agent.status === "queued", lines, budget, theme, () => totalRows++);
-    appendAgentGroup(phase.agents, (agent) => agent.status === "failed", lines, budget, theme, () => totalRows++);
-    appendAgentGroup(phase.agents, (agent) => agent.status === "done", lines, budget, theme, () => totalRows++);
+/**
+ * Spread the row budget over active phases by importance — a reserved share for failures so they always
+ * show, then running, remaining failures, done, queued — while each phase lists its chosen rows as
+ * running, failed, done, queued.
+ */
+function planAgentRows(active: readonly PhaseSnapshot[], budget: number): Map<PhaseSnapshot, AgentRowPlan> {
+  const byStatus = (status: AgentRowSnapshot["status"]) =>
+    active.flatMap((phase) => phase.agents.filter((agent) => agent.status === status));
+  const failed = byStatus("failed");
+  const chosen = new Set<AgentRowSnapshot>();
+  let remaining = budget;
+  const take = (agents: readonly AgentRowSnapshot[], limit = remaining) => {
+    for (const agent of agents) {
+      if (remaining <= 0 || limit <= 0) return;
+      if (chosen.has(agent)) continue;
+      chosen.add(agent);
+      remaining--;
+      limit--;
+    }
   };
+  take(failed, failed.length > 0 ? Math.max(1, Math.floor(budget / 3)) : 0);
+  take(byStatus("running"));
+  take(failed);
+  take(byStatus("done"));
+  take(byStatus("queued"));
 
-  for (const phase of phases) {
-    if (phase.agents.some((agent) => agent.status === "running" || agent.status === "queued")) visit(phase);
-  }
-  for (const phase of phases) {
-    if (!phase.agents.some((agent) => agent.status === "running" || agent.status === "queued")) visit(phase);
+  const rank = { running: 0, failed: 1, done: 2, queued: 3 } as const;
+  const plans = new Map<PhaseSnapshot, AgentRowPlan>();
+  for (const phase of active) {
+    const shown = phase.agents.filter((agent) => chosen.has(agent)).sort((a, b) => rank[a.status] - rank[b.status]);
+    plans.set(phase, { phase, shown, hidden: phase.agents.length - shown.length });
   }
+  return plans;
+}
 
-  return { lines, hidden: Math.max(0, totalRows - lines.length) };
+function agentLabelWidth(plans: ReadonlyMap<PhaseSnapshot, AgentRowPlan>, width: number): number {
+  const labels = [...plans.values()].flatMap((plan) => plan.shown.map((agent) => visibleWidth(agentLabel(agent, plan.phase))));
+  return Math.min(MAX_LABEL_COLUMN, Math.floor(width * 0.4), Math.max(0, ...labels));
 }
 
-function appendAgentGroup(
-  agents: readonly AgentRowSnapshot[],
-  include: (agent: AgentRowSnapshot) => boolean,
-  lines: string[],
-  budget: number,
-  theme: Theme,
-  countRow: () => void,
-): void {
-  for (const agent of agents) {
-    if (!include(agent)) continue;
-    countRow();
-    if (lines.length < budget) lines.push(agentLine(agent, theme));
+function finishedPhaseLine(phase: PhaseSnapshot, titleWidth: number, theme: Theme): string {
+  const counts = countAgents([phase]);
+  const glyph = counts.failed === 0
+    ? theme.fg("success", GLYPH.done)
+    : counts.failed === counts.total ? theme.fg("error", GLYPH.failed) : theme.fg("warning", GLYPH.warning);
+  const span = phaseSpan(phase.agents);
+  const detail = joinParts([
+    counts.total > 0 ? theme.fg("muted", plural(counts.total, "agent")) : undefined,
+    span === undefined ? undefined : theme.fg("muted", shortDuration(span)),
+    counts.failed > 0 ? theme.fg("error", `${counts.failed} failed`) : undefined,
+  ], theme);
+  return `  ${glyph} ${theme.fg("muted", fit(phase.title, titleWidth))}  ${detail}`.trimEnd();
+}
+
+function foldedPhasesLine(phases: readonly PhaseSnapshot[], theme: Theme): string {
+  const counts = countAgents(phases);
+  const glyph = counts.failed > 0 ? theme.fg("warning", GLYPH.warning) : theme.fg("success", GLYPH.done);
+  const titles = phases.map((phase) => phase.title).join(" · ");
+  const detail = joinParts([
+    theme.fg("muted", plural(counts.total, "agent")),
+    counts.failed > 0 ? theme.fg("error", `${counts.failed} failed`) : undefined,
+  ], theme);
+  return `  ${glyph} ${theme.fg("muted", titles)}  ${detail}`;
+}
+
+function activePhaseLine(plan: AgentRowPlan, titleWidth: number, theme: Theme): string {
+  const counts = countAgents([plan.phase]);
+  const detail = joinParts([
+    counts.running > 0 ? theme.fg("accent", `${counts.running} running`) : undefined,
+    counts.queued > 0 ? theme.fg("muted", `${counts.queued} queued`) : undefined,
+    counts.done > 0 ? theme.fg("muted", `${counts.done} done`) : undefined,
+    counts.failed > 0 ? theme.fg("error", `${counts.failed} failed`) : undefined,
+    plan.hidden > 0 ? theme.fg("dim", `+${plan.hidden} more`) : undefined,
+  ], theme);
+  return `  ${theme.fg("accent", GLYPH.running)} ${theme.bold(fit(plan.phase.title, titleWidth))}  ${detail}`.trimEnd();
+}
+
+function agentLine(agent: AgentRowSnapshot, phase: PhaseSnapshot, labelWidth: number, width: number, now: number, theme: Theme): string {
+  const glyph = agentGlyph(agent, theme);
+  const label = fit(agentLabel(agent, phase), labelWidth);
+  const head = `    ${glyph} ${theme.fg(agent.status === "running" || agent.status === "failed" ? "text" : "muted", label)}  `;
+  return head + truncateToWidth(agentDetail(agent, now, theme), Math.max(0, width - visibleWidth(head)), "…");
+}
+
+function agentGlyph(agent: AgentRowSnapshot, theme: Theme): string {
+  switch (agent.status) {
+    case "running":
+      return theme.fg("accent", GLYPH.running);
+    case "failed":
+      return theme.fg("error", GLYPH.failed);
+    case "done":
+      return theme.fg("success", GLYPH.done);
+    case "queued":
+      return theme.fg("dim", GLYPH.queued);
+  }
+}
+
+/** Running: what it is doing now; done: how much it did; failed: why. */
+function agentDetail(agent: AgentRowSnapshot, now: number, theme: Theme): string {
+  const elapsed = agent.startedAt === undefined ? undefined : (agent.doneAt ?? now) - agent.startedAt;
+  switch (agent.status) {
+    case "running":
+      return joinParts([agent.lastTool && theme.fg("muted", agent.lastTool), elapsed !== undefined && theme.fg("dim", formatElapsed(elapsed))], theme);
+    case "done":
+      return joinParts([
+        agent.toolUses > 0 && theme.fg("dim", plural(agent.toolUses, "tool")),
+        elapsed !== undefined && theme.fg("dim", shortDuration(elapsed)),
+      ], theme);
+    case "failed":
+      return theme.fg("error", oneLine(agent.error ?? "failed"));
+    case "queued":
+      return theme.fg("dim", "queued");
   }
 }
 
-function countPhaseAgents(phase: PhaseSnapshot): AgentCounts {
-  const counts: AgentCounts = { queued: 0, running: 0, done: 0, failed: 0, total: phase.agents.length };
-  for (const agent of phase.agents) counts[agent.status]++;
-  return counts;
+/** `find:logic-bugs` under the Find phase reads as `logic-bugs`; the phase line already says Find. */
+function agentLabel(agent: AgentRowSnapshot, phase: PhaseSnapshot): string {
+  const phaseName = phase.title.split(" ▸ ").pop()?.toLowerCase() ?? "";
+  const prefix = `${phaseName}:`;
+  return phaseName && agent.label.toLowerCase().startsWith(prefix) && agent.label.length > prefix.length
+    ? agent.label.slice(prefix.length)
+    : agent.label;
+}
+
+/** Declared `meta.phases` not reached yet, after the last one the run has entered. */
+function upcomingLine(snapshot: WorkflowProgressSnapshot, theme: Theme): string | undefined {
+  const planned = snapshot.plannedPhases ?? [];
+  const seen = new Set(snapshot.phases.map((phase) => phase.title));
+  const lastSeen = planned.reduce((last, title, index) => (seen.has(title) ? index : last), -1);
+  const upcoming = planned.slice(lastSeen + 1).filter((title) => !seen.has(title));
+  if (upcoming.length === 0) return undefined;
+  return `  ${theme.fg("dim", GLYPH.queued)} ${theme.fg("dim", upcoming.join(" · "))}`;
 }
 
-function phaseLine(title: string, counts: AgentCounts, theme: Theme): string {
+/**
+ * The latest log when it reports a failure the phase and agent lines do not already count, such as a
+ * provider retry or an observer error. Failed agents are reported by their rows and phase totals.
+ */
+function failureLogLine(snapshot: WorkflowProgressSnapshot, theme: Theme): string | undefined {
+  const latest = snapshot.logs.at(-1);
+  if (!latest || !FAILURE_LOG.test(latest)) return undefined;
+  const failedLabels = snapshot.phases.flatMap((phase) => phase.agents.filter((agent) => agent.status === "failed").map((agent) => agent.label));
+  if (failedLabels.some((label) => latest.startsWith(label))) return undefined;
+  return `  ${theme.fg("warning", `${GLYPH.warning} ${oneLine(latest)}`)}`;
+}
+
+function footerLine(snapshot: WorkflowProgressSnapshot, inspectorHint: string, width: number, theme: Theme): string {
+  const hint = keyHints([[inspectorHint, "inspect"]], theme);
+  const counters = counterText(snapshot, theme);
+  return counters ? spread(`  ${counters}`, hint, width) : spread("", hint, width);
+}
+
+/** `3 files · 3 candidates · 2 confirmed`: zeros hidden, verdict labels lower-cased, duplicates dropped. */
+function counterText(snapshot: WorkflowProgressSnapshot, theme: Theme): string | undefined {
+  const seen = new Set<string>();
   const parts: string[] = [];
-  if (counts.running > 0) parts.push(`${counts.running} running`);
-  if (counts.done > 0) parts.push(`${counts.done} done`);
-  if (counts.failed > 0) parts.push(`${counts.failed} failed`);
-  return `${theme.fg("dim", "├─")} ${theme.fg(counts.running > 0 ? "accent" : "muted", title)}${parts.length ? ` ${theme.fg("dim", `(${parts.join(" · ")})`)}` : ""}`;
-}
-
-function agentLine(agent: AgentRowSnapshot, theme: Theme): string {
-  const detailParts = agentDetailParts(agent, { includeQueuedStatus: false });
-  const activity = detailParts.length > 0 ? ` ${theme.fg("dim", `· ${detailParts.join(" · ")}`)}` : "";
-  return `${theme.fg("dim", "│  ")} ${statusIcon(agent.status, theme)} ${theme.fg(agentLabelColor(agent), agent.label)}${activity}`;
-}
-
-function footerLine(snapshot: WorkflowProgressSnapshot, theme: Theme): string | undefined {
-  const usage = formatWorkflowUsageLine(snapshot.usage);
-  const counters = snapshot.counters.slice(0, 4).map((counter) => `${counter.label} ${formatCount(counter.value)}`);
-  const latestLog = snapshot.logs[snapshot.logs.length - 1];
-  const parts = usage ? [usage, ...counters] : counters;
-  if (latestLog) parts.push(latestLog);
-  if (parts.length === 0) return undefined;
-  return `${theme.fg("dim", "└─")} ${theme.fg("dim", parts.join(" · "))}`;
+  for (const counter of snapshot.counters) {
+    const label = counter.label.toLowerCase().replaceAll("_", " ");
+    if (counter.value === 0 || seen.has(label)) continue;
+    seen.add(label);
+    parts.push(`${theme.fg("text", formatCount(counter.value))} ${theme.fg(counterColor(counter.key), label)}`);
+  }
+  return parts.length > 0 ? parts.join(dot(theme)) : undefined;
+}
+
+function counterColor(key: string): ThemeColor {
+  return key === "kept" ? "success" : "muted";
+}
+
+function phaseSpan(agents: readonly AgentRowSnapshot[]): number | undefined {
+  const starts = agents.flatMap((agent) => (agent.startedAt === undefined ? [] : [agent.startedAt]));
+  const ends = agents.flatMap((agent) => (agent.doneAt === undefined ? [] : [agent.doneAt]));
+  if (starts.length === 0 || ends.length === 0) return undefined;
+  return Math.max(...ends) - Math.min(...starts);
+}
+
+/** `0.6s`, `4.2s`, then `12s` / `1m 4s` for finished spans. */
+function shortDuration(ms: number): string {
+  if (ms < 10_000) return `${(Math.max(0, ms) / 1_000).toFixed(1)}s`;
+  return formatDuration(ms);
+}
+
+function plural(count: number, noun: string): string {
+  return `${count} ${noun}${count === 1 ? "" : "s"}`;
+}
+
+function oneLine(text: string): string {
+  return text.replace(/\s+/g, " ").trim();
 }

package/.pi/extensions/pi-workflow-engine/src/usage.ts

@@ -1,3 +1,7 @@
+import type { Usage } from "@earendil-works/pi-ai";
+import { isFiniteNumber, isRecord } from "./guards.ts";
+import { formatCount } from "./text.ts";
+
 export interface WorkflowUsageCost {
   readonly input: number;
   readonly output: number;
@@ -78,9 +82,14 @@ export class WorkflowUsageRecorder implements WorkflowUsageSink {
       const parsed = parseAssistantUsageMessage(message);
       return parsed ? [parsed] : [];
     });
-    if (assistantMessages.length === 0) return;
+    // Tools that call models themselves, such as codemode running a classifier, report it on their result.
+    const toolUsage = input.messages.flatMap((message) => {
+      const parsed = parseToolResultUsage(message);
+      return parsed ? [parsed] : [];
+    });
+    if (assistantMessages.length === 0 && toolUsage.length === 0) return;
 
-    const usage = sumTotals(assistantMessages.map((message) => message.usage));
+    const usage = sumTotals([...assistantMessages.map((message) => message.usage), ...toolUsage]);
     const latestMetadata = assistantMessages.findLast((message) => message.provider !== undefined || message.model !== undefined);
     this.agents.push({
       label: input.label,
@@ -94,10 +103,8 @@ export class WorkflowUsageRecorder implements WorkflowUsageSink {
   }
 
   snapshot(): WorkflowUsageSnapshot {
-    const agents = this.agents.map((agent) => ({
-      ...agent,
-      usage: cloneTotals(agent.usage),
-    }));
+    // Recorded entries are never mutated; copying the array keeps each snapshot point-in-time.
+    const agents = [...this.agents];
     return {
       agents,
       totals: sumTotals(agents.map((agent) => agent.usage)),
@@ -118,7 +125,7 @@ export function isWorkflowUsageSnapshot(value: unknown): value is WorkflowUsageS
   if (!isRecord(value)) return false;
   if (!Array.isArray(value.agents) || !value.agents.every(isWorkflowAgentUsage)) return false;
   if (!isWorkflowUsageTotals(value.totals)) return false;
-  return finiteNumber(value.assistantMessages) !== undefined;
+  return isFiniteNumber(value.assistantMessages);
 }
 
 export function hasWorkflowUsage(snapshot: unknown): snapshot is WorkflowUsageSnapshot {
@@ -126,6 +133,13 @@ export function hasWorkflowUsage(snapshot: unknown): snapshot is WorkflowUsageSn
   return snapshot.totals.totalTokens > 0 || snapshot.totals.cost.total > 0;
 }
 
+/** The run's subagent totals as pi's `Usage`, so a tool result adds them to the session's cost. */
+export function workflowUsageForPi(snapshot: unknown): Usage | undefined {
+  if (!hasWorkflowUsage(snapshot)) return undefined;
+  const { input, output, cacheRead, cacheWrite, totalTokens, cost } = snapshot.totals;
+  return { input, output, cacheRead, cacheWrite, totalTokens, cost: { ...cost } };
+}
+
 export function formatWorkflowUsageLine(snapshot: unknown): string | undefined {
   if (!hasWorkflowUsage(snapshot)) return undefined;
   const parts = formatWorkflowUsageComponents(snapshot.totals);
@@ -144,9 +158,9 @@ function formatWorkflowUsageComponents(totals: WorkflowUsageTotals): string[] {
   appendUsageComponent(parts, "output", totals.output, totals.coverage.output);
 
   const knownComponents = totals.input + totals.output + totals.cacheRead + totals.cacheWrite;
-  if (parts.length === 0 && totals.totalTokens > 0) parts.push(`tokens ${formatUsageCount(totals.totalTokens)}`);
+  if (parts.length === 0 && totals.totalTokens > 0) parts.push(`tokens ${formatCount(totals.totalTokens)}`);
   else if (!hasCompleteCoverage(totals.coverage) && totals.totalTokens !== knownComponents) {
-    parts.push(`total ${formatUsageCount(totals.totalTokens)}`);
+    parts.push(`total ${formatCount(totals.totalTokens)}`);
   }
   return parts;
 }
@@ -158,7 +172,7 @@ function appendUsageComponent(
   coverage: WorkflowUsageCoverage,
 ): void {
   if (value <= 0 || coverage === "none") return;
-  const count = formatUsageCount(value);
+  const count = formatCount(value);
   parts.push(`${label} ${coverage === "partial" ? `≥${count}` : count}`);
 }
 
@@ -168,17 +182,17 @@ function isWorkflowAgentUsage(value: unknown): value is WorkflowAgentUsage {
   if (value.phase !== undefined && typeof value.phase !== "string") return false;
   if (value.provider !== undefined && typeof value.provider !== "string") return false;
   if (value.model !== undefined && typeof value.model !== "string") return false;
-  if (finiteNumber(value.assistantMessages) === undefined) return false;
+  if (!isFiniteNumber(value.assistantMessages)) return false;
   return isWorkflowUsageTotals(value.usage);
 }
 
 function isWorkflowUsageTotals(value: unknown): value is WorkflowUsageTotals {
   if (!isRecord(value)) return false;
-  if (finiteNumber(value.input) === undefined) return false;
-  if (finiteNumber(value.output) === undefined) return false;
-  if (finiteNumber(value.cacheRead) === undefined) return false;
-  if (finiteNumber(value.cacheWrite) === undefined) return false;
-  if (finiteNumber(value.totalTokens) === undefined) return false;
+  if (!isFiniteNumber(value.input)) return false;
+  if (!isFiniteNumber(value.output)) return false;
+  if (!isFiniteNumber(value.cacheRead)) return false;
+  if (!isFiniteNumber(value.cacheWrite)) return false;
+  if (!isFiniteNumber(value.totalTokens)) return false;
   if (!isWorkflowUsageCoverage(value.coverage)) return false;
   return isWorkflowUsageCost(value.cost);
 }
@@ -196,11 +210,11 @@ function isWorkflowUsageCoverage(value: unknown): value is WorkflowUsageComponen
 function isWorkflowUsageCost(value: unknown): value is WorkflowUsageCost {
   if (!isRecord(value)) return false;
   return (
-    finiteNumber(value.input) !== undefined &&
-    finiteNumber(value.output) !== undefined &&
-    finiteNumber(value.cacheRead) !== undefined &&
-    finiteNumber(value.cacheWrite) !== undefined &&
-    finiteNumber(value.total) !== undefined
+    isFiniteNumber(value.input) &&
+    isFiniteNumber(value.output) &&
+    isFiniteNumber(value.cacheRead) &&
+    isFiniteNumber(value.cacheWrite) &&
+    isFiniteNumber(value.total)
   );
 }
 
@@ -216,6 +230,11 @@ function parseAssistantUsageMessage(message: unknown): AssistantUsageMessage | u
   };
 }
 
+function parseToolResultUsage(message: unknown): WorkflowUsageTotals | undefined {
+  if (!isRecord(message) || message.role !== "toolResult") return undefined;
+  return parseUsageTotals(message.usage);
+}
+
 function parseUsageTotals(value: unknown): WorkflowUsageTotals | undefined {
   if (!isRecord(value)) return undefined;
   const input = finiteNumber(value.input);
@@ -321,18 +340,6 @@ function isCoverageValue(value: unknown): value is WorkflowUsageCoverage {
   return value === "none" || value === "partial" || value === "complete";
 }
 
-function formatUsageCount(count: number): string {
-  if (count < 1000) return count.toString();
-  if (count < 10000) return `${(count / 1000).toFixed(1)}k`;
-  if (count < 1000000) return `${Math.round(count / 1000)}k`;
-  if (count < 10000000) return `${(count / 1000000).toFixed(1)}M`;
-  return `${Math.round(count / 1000000)}M`;
-}
-
 function finiteNumber(value: unknown): number | undefined {
-  return typeof value === "number" && Number.isFinite(value) ? value : undefined;
-}
-
-function isRecord(value: unknown): value is Record<string, unknown> {
-  return typeof value === "object" && value !== null;
+  return isFiniteNumber(value) ? value : undefined;
 }