@aaroncql/pim-agent 0.0.1

package/src/extensions/subagent/render.ts

@@ -0,0 +1,359 @@
+import type {
+  AgentToolResult,
+  Theme,
+  ThemeColor,
+  ToolRenderResultOptions,
+} from "@earendil-works/pi-coding-agent";
+import type {
+  Component,
+  DefaultTextStyle,
+  MarkdownTheme,
+} from "@earendil-works/pi-tui";
+import { Container, Markdown, visibleWidth } from "@earendil-works/pi-tui";
+import { type PrefixSpec, Renderer } from "../../shared/Renderer";
+import type { SubagentDetails, SubagentSnapshot } from "./subagent";
+
+const DOT = "⬝";
+
+type RenderContext = {
+  readonly lastComponent: Component | undefined;
+  readonly isPartial: boolean;
+  readonly isError: boolean;
+};
+
+type StatusFields = Pick<
+  SubagentSnapshot,
+  | "usage"
+  | "toolCalls"
+  | "activeToolNames"
+  | "lastToolName"
+  | "stopReason"
+  | "model"
+  | "contextWindow"
+>;
+
+type MarkdownBlockArgs = {
+  readonly text: string;
+  readonly theme: Theme;
+  readonly prefix: PrefixSpec;
+  readonly lineColor?: ThemeColor;
+};
+
+class MarkdownTitle implements Component {
+  private label = "";
+  private title = "";
+  private theme: Theme | undefined;
+  private context: RenderContext | undefined;
+  private labelColor: ThemeColor | undefined;
+
+  public set(args: {
+    readonly label: string;
+    readonly title: string;
+    readonly theme: Theme;
+    readonly context: RenderContext;
+    readonly labelColor?: ThemeColor;
+  }): void {
+    this.label = args.label;
+    this.title = args.title;
+    this.theme = args.theme;
+    this.context = args.context;
+    this.labelColor = args.labelColor;
+  }
+
+  public render(width: number): string[] {
+    const theme = this.theme;
+    const context = this.context;
+    if (!theme || !context) {
+      return [];
+    }
+
+    const markerColor = Renderer.markerColorFor(
+      Boolean(context.isPartial),
+      Boolean(context.isError)
+    );
+    const prefix =
+      theme.fg(markerColor, " ▪") +
+      " " +
+      theme.fg(this.labelColor ?? "toolTitle", theme.bold(this.label)) +
+      theme.fg("toolTitle", ": ");
+    const inner = Math.max(1, width - visibleWidth(prefix));
+    const titleLines = renderMarkdownLines({
+      text: this.title,
+      theme,
+      width: inner,
+    });
+    const lines = titleLines.length > 0 ? titleLines : [""];
+    const out = [padLine(prefix + (lines[0] ?? ""), width)];
+
+    for (const line of lines.slice(1)) {
+      out.push(
+        padLine(
+          theme.fg("toolOutput", Renderer.GAPPED_PREFIX.prefix) + line,
+          width
+        )
+      );
+    }
+
+    return out;
+  }
+
+  public invalidate(): void {}
+}
+
+export function formatCallTitle(prompt: string | undefined): string {
+  return (prompt ?? "...").split(/\r?\n/u)[0]?.trim() || "...";
+}
+
+export function formatTopLine(snapshot: StatusFields): string {
+  return [
+    formatCost(snapshot.usage.cost),
+    formatContext(snapshot),
+    snapshot.model ?? "unknown model",
+    formatActivity(snapshot),
+  ].join(` ${DOT} `);
+}
+
+export function renderCall(
+  args: { readonly prompt?: string } | undefined,
+  theme: Theme,
+  context: RenderContext
+): Component {
+  const component =
+    context.lastComponent instanceof MarkdownTitle
+      ? context.lastComponent
+      : new MarkdownTitle();
+  component.set({
+    label: "Subagent",
+    title: formatCallTitle(args?.prompt),
+    theme,
+    context,
+    labelColor: titleColorFor(context),
+  });
+  return component;
+}
+
+export function renderResult(
+  result: AgentToolResult<SubagentDetails>,
+  options: ToolRenderResultOptions,
+  theme: Theme,
+  context: RenderContext
+): Component {
+  const container =
+    context.lastComponent instanceof Container
+      ? context.lastComponent
+      : new Container();
+  container.clear();
+
+  const details = result.details;
+  const first = result.content?.[0];
+  const body = first && "text" in first ? (first.text ?? "") : "";
+  const topLine = details?.topLine;
+
+  if (topLine) {
+    container.addChild(
+      Renderer.makePrefixedBlock({
+        text: styleDottedLine({
+          text: topLine,
+          theme,
+          lineColor: options.isPartial ? "warning" : "accent",
+        }),
+        theme,
+        prefix: Renderer.GAPPED_PREFIX,
+      })
+    );
+  }
+
+  if (!options.isPartial && (!topLine || options.expanded) && body) {
+    container.addChild(
+      options.expanded
+        ? makePrefixedMarkdownBlock({
+            text: body,
+            theme,
+            prefix: Renderer.GAPPED_PREFIX,
+            lineColor: expandedResultColor(details, context.isError),
+          })
+        : Renderer.makePrefixedBlock({
+            text: body,
+            theme,
+            prefix: Renderer.GAPPED_PREFIX,
+            lineColor: resultColor(details, context.isError),
+          })
+    );
+  }
+
+  container.invalidate();
+  return container;
+}
+
+function makePrefixedMarkdownBlock(args: MarkdownBlockArgs): Component {
+  const markdown = new Markdown(
+    args.text,
+    0,
+    0,
+    makeMarkdownTheme(args.theme),
+    defaultStyle(args.theme, args.lineColor)
+  );
+
+  return {
+    render(width: number): string[] {
+      const inner = Math.max(1, width - args.prefix.width);
+      return markdown.render(inner).map((line) => {
+        return (
+          args.theme.fg("toolOutput", args.prefix.prefix) +
+          trimRenderedLine(line)
+        );
+      });
+    },
+    invalidate(): void {
+      markdown.invalidate();
+    },
+  };
+}
+
+function renderMarkdownLines(args: {
+  readonly text: string;
+  readonly theme: Theme;
+  readonly width: number;
+  readonly lineColor?: ThemeColor;
+}): string[] {
+  const markdown = new Markdown(
+    args.text,
+    0,
+    0,
+    makeMarkdownTheme(args.theme),
+    defaultStyle(args.theme, args.lineColor)
+  );
+  return markdown.render(args.width).map(trimRenderedLine);
+}
+
+function defaultStyle(
+  theme: Theme,
+  lineColor: ThemeColor | undefined
+): DefaultTextStyle | undefined {
+  return lineColor
+    ? { color: (text: string) => theme.fg(lineColor, text) }
+    : undefined;
+}
+
+function makeMarkdownTheme(theme: Theme): MarkdownTheme {
+  return {
+    heading: (text: string) => theme.fg("mdHeading", text),
+    link: (text: string) => theme.fg("mdLink", text),
+    linkUrl: (text: string) => theme.fg("mdLinkUrl", text),
+    code: (text: string) => theme.fg("mdCode", text),
+    codeBlock: (text: string) => theme.fg("mdCodeBlock", text),
+    codeBlockBorder: (text: string) => theme.fg("mdCodeBlockBorder", text),
+    quote: (text: string) => theme.fg("mdQuote", text),
+    quoteBorder: (text: string) => theme.fg("mdQuoteBorder", text),
+    hr: (text: string) => theme.fg("mdHr", text),
+    listBullet: (text: string) => theme.fg("mdListBullet", text),
+    bold: (text: string) => theme.bold(text),
+    italic: (text: string) => theme.italic(text),
+    underline: (text: string) => theme.underline(text),
+    strikethrough: (text: string) => theme.strikethrough(text),
+  };
+}
+
+function trimRenderedLine(line: string): string {
+  return line.trimEnd();
+}
+
+function padLine(line: string, width: number): string {
+  return line + " ".repeat(Math.max(0, width - visibleWidth(line)));
+}
+
+function styleDottedLine(args: {
+  readonly text: string;
+  readonly theme: Theme;
+  readonly lineColor: ThemeColor;
+}): string {
+  return args.text
+    .split(DOT)
+    .map((part) => args.theme.fg(args.lineColor, part))
+    .join(args.theme.fg("muted", DOT));
+}
+
+function titleColorFor(context: RenderContext): ThemeColor {
+  if (context.isPartial) {
+    return "warning";
+  }
+  if (context.isError) {
+    return "error";
+  }
+  return "accent";
+}
+
+function resultColor(
+  details: SubagentDetails | undefined,
+  isError: boolean
+): "toolOutput" | "error" | "warning" {
+  if (details?.stopReason === "aborted") {
+    return "warning";
+  }
+  return isError ? "error" : "toolOutput";
+}
+
+function expandedResultColor(
+  details: SubagentDetails | undefined,
+  isError: boolean
+): "error" | "warning" | undefined {
+  if (details?.stopReason === "aborted") {
+    return "warning";
+  }
+  return isError ? "error" : undefined;
+}
+
+function formatActivity(snapshot: StatusFields): string {
+  const turns = `${snapshot.usage.turns} ${snapshot.usage.turns === 1 ? "turn" : "turns"}`;
+  if (snapshot.stopReason !== undefined) {
+    const toolCount = snapshot.toolCalls.length;
+    return toolCount > 0
+      ? `${turns} ${DOT} ${toolCount} ${toolCount === 1 ? "tool" : "tools"}`
+      : turns;
+  }
+
+  return `${turns} ${DOT} ${activeToolLabel(snapshot)}`;
+}
+
+function activeToolLabel(snapshot: StatusFields): string {
+  if (snapshot.activeToolNames.length === 1) {
+    return snapshot.activeToolNames[0]!;
+  }
+  if (snapshot.activeToolNames.length > 1) {
+    return `${snapshot.activeToolNames.length} tools`;
+  }
+  return snapshot.lastToolName ?? "thinking";
+}
+
+function formatContext(snapshot: StatusFields): string {
+  const window = snapshot.contextWindow;
+  if (!window || window <= 0) {
+    return "?/?";
+  }
+  const windowText = formatTokens(window);
+  const tokens = snapshot.usage.contextTokens;
+  if (tokens === undefined) {
+    return `?/${windowText}`;
+  }
+  return `${((tokens / window) * 100).toFixed(1)}%/${windowText}`;
+}
+
+function formatTokens(tokens: number): string {
+  if (tokens < 1000) {
+    return `${tokens}`;
+  }
+  if (tokens < 10_000) {
+    return `${(tokens / 1000).toFixed(1)}K`;
+  }
+  if (tokens < 1_000_000) {
+    return `${Math.round(tokens / 1000)}K`;
+  }
+  if (tokens < 10_000_000) {
+    return `${(tokens / 1_000_000).toFixed(1)}M`;
+  }
+  return `${Math.round(tokens / 1_000_000)}M`;
+}
+
+function formatCost(cost: number): string {
+  return `$${cost.toFixed(2)}`;
+}

package/src/extensions/subagent/schema.ts

@@ -0,0 +1,9 @@
+import { type Static, Type } from "typebox";
+
+export const subagentSchema = Type.Object({
+  prompt: Type.String({
+    minLength: 1,
+  }),
+});
+
+export type SubagentInput = Static<typeof subagentSchema>;

package/src/extensions/subagent/subagent.test.ts

@@ -0,0 +1,315 @@
+import { describe, expect, test } from "bun:test";
+import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
+import type { AssistantMessage, Usage } from "@earendil-works/pi-ai";
+import {
+  applyOutputCap,
+  childToolNames,
+  runSubagent,
+  SubagentEventCapture,
+  type SubagentSession,
+} from "./subagent";
+
+type UsageOverrides = Omit<Partial<Usage>, "cost"> & {
+  readonly cost?: Partial<Usage["cost"]>;
+};
+
+const usage = (overrides: UsageOverrides = {}): Usage => {
+  const { cost, ...rest } = overrides;
+  return {
+    input: 0,
+    output: 0,
+    cacheRead: 0,
+    cacheWrite: 0,
+    totalTokens: 0,
+    cost: {
+      input: 0,
+      output: 0,
+      cacheRead: 0,
+      cacheWrite: 0,
+      total: 0,
+      ...cost,
+    },
+    ...rest,
+  };
+};
+
+function assistant(
+  textParts: readonly string[],
+  overrides: Partial<AssistantMessage> = {}
+): AssistantMessage {
+  return {
+    role: "assistant",
+    content: textParts.map((text) => ({ type: "text", text })),
+    api: "anthropic-messages",
+    provider: "anthropic",
+    model: "claude-test",
+    usage: usage(),
+    stopReason: "stop",
+    timestamp: 1,
+    ...overrides,
+  };
+}
+
+const ctx = { cwd: "/work" } as ExtensionContext;
+
+class FakeSession implements SubagentSession {
+  public promptCalls = 0;
+  public abortCalls = 0;
+  public disposeCalls = 0;
+  private listener: ((event: never) => void) | undefined;
+
+  public constructor(
+    private readonly onPrompt: (
+      session: FakeSession,
+      prompt: string
+    ) => Promise<void>
+  ) {}
+
+  public subscribe(listener: (event: never) => void): () => void {
+    this.listener = listener;
+    return () => {
+      this.listener = undefined;
+    };
+  }
+
+  public emit(event: unknown): void {
+    this.listener?.(event as never);
+  }
+
+  public async prompt(prompt: string): Promise<void> {
+    this.promptCalls += 1;
+    await this.onPrompt(this, prompt);
+  }
+
+  public async abort(): Promise<void> {
+    this.abortCalls += 1;
+  }
+
+  public dispose(): void {
+    this.disposeCalls += 1;
+  }
+}
+
+describe("childToolNames", () => {
+  test("removes the subagent tool from a child's inherited allowlist", () => {
+    expect(childToolNames(["read", "subagent", "bash"])).toEqual([
+      "read",
+      "bash",
+    ]);
+  });
+});
+
+describe("SubagentEventCapture", () => {
+  test("concatenates multi-part text, resets for each assistant message, and records usage/tools", () => {
+    const updates: string[] = [];
+    const capture = new SubagentEventCapture((partial) => {
+      updates.push(
+        partial.content[0]?.type === "text" ? partial.content[0].text : ""
+      );
+    });
+
+    capture.handle({ type: "message_start", message: assistant([]) } as never);
+    capture.handle({
+      type: "message_end",
+      message: assistant(["first ", "turn"], {
+        usage: usage({ input: 10, output: 4, cost: { total: 0.01 } }),
+      }),
+    } as never);
+    capture.handle({
+      type: "tool_execution_end",
+      toolCallId: "1",
+      toolName: "read",
+      result: {},
+      isError: false,
+    } as never);
+    capture.handle({ type: "message_start", message: assistant([]) } as never);
+    capture.handle({
+      type: "message_end",
+      message: assistant(["final", " answer"], {
+        usage: usage({
+          input: 2,
+          output: 8,
+          cacheRead: 3,
+          cost: { total: 0.02 },
+        }),
+      }),
+    } as never);
+
+    const snapshot = capture.snapshot();
+    expect(snapshot.finalOutput).toBe("final answer");
+    expect(snapshot.usage).toEqual({
+      input: 12,
+      output: 12,
+      cacheRead: 3,
+      cacheWrite: 0,
+      cost: 0.03,
+      turns: 2,
+      contextTokens: undefined,
+    });
+    expect(snapshot.toolCalls).toEqual([{ name: "read", isError: false }]);
+    expect(snapshot.lastToolName).toBe("read");
+    expect(updates.at(-1)).toBe("$0.03 ⬝ ?/? ⬝ claude-test ⬝ 2 turns ⬝ 1 tool");
+  });
+
+  test("a later message with no text discards the prior message's text", () => {
+    const capture = new SubagentEventCapture();
+
+    capture.handle({ type: "message_start", message: assistant([]) } as never);
+    capture.handle({
+      type: "message_end",
+      message: assistant(["intro"], { stopReason: "toolUse" }),
+    } as never);
+    capture.handle({ type: "message_start", message: assistant([]) } as never);
+    capture.handle({
+      type: "message_end",
+      message: assistant([], { stopReason: "stop" }),
+    } as never);
+
+    expect(capture.snapshot().finalOutput).toBe("");
+  });
+
+  test("message_update content is materialized lazily on snapshot read", () => {
+    const capture = new SubagentEventCapture();
+
+    capture.handle({ type: "message_start", message: assistant([]) } as never);
+    capture.handle({
+      type: "message_update",
+      message: assistant(["partial"]),
+    } as never);
+
+    expect(capture.snapshot().finalOutput).toBe("partial");
+  });
+});
+
+describe("applyOutputCap", () => {
+  test("truncates on a UTF-8 boundary and reports omitted bytes", () => {
+    const capped = applyOutputCap("😀😀😀", 5);
+
+    expect(capped.text).toContain(
+      "😀\n[subagent: output truncated, 8 bytes omitted"
+    );
+    expect(capped.text).not.toContain("�");
+    expect(capped.truncated).toBe(true);
+    expect(capped.omittedBytes).toBe(8);
+  });
+});
+
+describe("runSubagent", () => {
+  test("returns bare text on normal completion", async () => {
+    const fake = new FakeSession(async (session) => {
+      session.emit({ type: "message_start", message: assistant([]) });
+      session.emit({
+        type: "message_end",
+        message: assistant(["hello"]),
+      });
+    });
+
+    const result = await runSubagent(
+      "say hi",
+      ctx,
+      undefined,
+      undefined,
+      async () => fake
+    );
+
+    expect(result.content).toEqual([{ type: "text", text: "hello" }]);
+    expect(result.details.fullOutput).toBe("hello");
+    expect(fake.promptCalls).toBe(1);
+    expect(fake.abortCalls).toBe(1);
+    expect(fake.disposeCalls).toBe(1);
+  });
+
+  test("returns a hint, not an error, for normal empty output", async () => {
+    const fake = new FakeSession(async (session) => {
+      session.emit({ type: "message_start", message: assistant([]) });
+      session.emit({ type: "message_end", message: assistant([]) });
+    });
+
+    const result = await runSubagent(
+      "empty",
+      ctx,
+      undefined,
+      undefined,
+      async () => fake
+    );
+
+    expect(
+      result.content[0]?.type === "text" ? result.content[0].text : ""
+    ).toBe("[subagent tool: completed with no text output.]");
+  });
+
+  test("throws on model error with partial output", async () => {
+    const fake = new FakeSession(async (session) => {
+      session.emit({ type: "message_start", message: assistant([]) });
+      session.emit({
+        type: "message_end",
+        message: assistant(["partial"], {
+          stopReason: "error",
+          errorMessage: "provider exploded",
+        }),
+      });
+    });
+
+    await expect(
+      runSubagent("fail", ctx, undefined, undefined, async () => fake)
+    ).rejects.toThrow(
+      "Subagent failed: error. Error: provider exploded.\nPartial output before failure:\npartial"
+    );
+  });
+
+  test("rejects pre-aborted signals before prompt", async () => {
+    const controller = new AbortController();
+    controller.abort();
+    const fake = new FakeSession(async () => {});
+
+    await expect(
+      runSubagent("abort", ctx, controller.signal, undefined, async () => fake)
+    ).rejects.toThrow("Subagent failed: subagent aborted before start");
+
+    expect(fake.promptCalls).toBe(0);
+    expect(fake.abortCalls).toBe(1);
+    expect(fake.disposeCalls).toBe(1);
+  });
+
+  test("mid-run abort aborts once, tears down, and rejects", async () => {
+    let finishPrompt: (() => void) | undefined;
+    const fake = new (class extends FakeSession {
+      public override async abort(): Promise<void> {
+        await super.abort();
+        finishPrompt?.();
+      }
+    })(
+      () =>
+        new Promise<void>((resolve) => {
+          finishPrompt = resolve;
+        })
+    );
+    const controller = new AbortController();
+    const promise = runSubagent(
+      "long",
+      ctx,
+      controller.signal,
+      undefined,
+      async () => fake
+    );
+
+    await Promise.resolve();
+    controller.abort();
+
+    await expect(promise).rejects.toThrow("Subagent failed: aborted");
+    expect(fake.promptCalls).toBe(1);
+    expect(fake.abortCalls).toBe(1);
+    expect(fake.disposeCalls).toBe(1);
+  });
+
+  test("nested subagent calls are rejected by the async-local recursion ban", async () => {
+    const outer = new FakeSession(async () => {
+      const inner = new FakeSession(async () => {});
+      await runSubagent("inner", ctx, undefined, undefined, async () => inner);
+    });
+
+    await expect(
+      runSubagent("outer", ctx, undefined, undefined, async () => outer)
+    ).rejects.toThrow("subagents cannot call subagent tool");
+  });
+});