feat: enable permission-aware subagents

pi-subagents/test/cost-display.test.ts

@@ -0,0 +1,217 @@
+/**
+ * cost-display.test.ts — what `showCost` puts on each surface, and what it must
+ * never put there.
+ *
+ * Two rules run through all of it. A cost is shown only when there is one to
+ * show: a model pi has no rates for reports 0, and printing `$0.00` beside its
+ * tokens would claim the run was measured and free. And each surface punctuates
+ * its own — the stats line joins with `·`, the foreground result with `,`, the
+ * `get_subagent_result` header with `|` — which is why the cost travels as a
+ * number and is formatted at the end, not baked into the token string.
+ */
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+vi.mock("../src/agent-runner.js", async () => {
+  const actual = await vi.importActual<typeof import("../src/agent-runner.js")>("../src/agent-runner.js");
+  return { ...actual, runAgent: vi.fn() };
+});
+
+import { runAgent } from "../src/agent-runner.js";
+import { registerAgents } from "../src/agent-types.js";
+import subagentsExtension from "../src/index.js";
+import { ctx, flush, type Hermetic, hermeticDir, makePi, textOf } from "./helpers/boot-extension.js";
+
+const COST = 0.0123;
+
+/** One foreground run that spends `cost` on a single assistant message. */
+function runSpending(cost: number) {
+  vi.mocked(runAgent).mockImplementation(async (_c: any, _t: any, _p: any, opts: any) => {
+    opts.onAssistantUsage?.({ input: 1000, output: 200, cacheWrite: 0, cost });
+    return { responseText: "done", session: { dispose: vi.fn() } as any, aborted: false, steered: false };
+  });
+}
+
+const spawn = (tools: Map<string, any>) =>
+  tools.get("Agent").execute(
+    "tc-1",
+    { prompt: "go", description: "spend", subagent_type: "general-purpose", run_in_background: false },
+    undefined, undefined, ctx(),
+  );
+
+describe("cost display", () => {
+  let hermetic: Hermetic;
+
+  function boot(settings: Record<string, unknown>) {
+    hermetic = hermeticDir({ settings });
+    const { pi, tools, lifecycle } = makePi();
+    subagentsExtension(pi);
+    return { pi, tools, lifecycle };
+  }
+
+  beforeEach(() => {
+    vi.mocked(runAgent).mockReset();
+  });
+
+  afterEach(() => {
+    delete (globalThis as any)[Symbol.for("pi-subagents:manager")];
+    registerAgents(new Map());
+    hermetic?.restore();
+  });
+
+  describe("the foreground result the orchestrator reads", () => {
+    it("names the cost in the stats it already reports", async () => {
+      const { tools } = boot({ showCost: true });
+      runSpending(COST);
+
+      const text = textOf(await spawn(tools));
+
+      expect(text).toContain("~$0.0123");
+      // Comma-joined with the rest, not glued to the token count with the "·"
+      // that the widget uses — the separator belongs to the surface.
+      expect(text).toMatch(/1\.2k token, ~\$0\.0123/);
+    });
+
+    it("says nothing when the setting is off", async () => {
+      const { tools } = boot({ showCost: false });
+      runSpending(COST);
+
+      expect(textOf(await spawn(tools))).not.toContain("$");
+    });
+
+    it("says nothing for a model with no pricing data", async () => {
+      const { tools } = boot({ showCost: true });
+      runSpending(0);
+
+      const text = textOf(await spawn(tools));
+      expect(text).toContain("1.2k token");   // tokens are still exact
+      expect(text).not.toContain("$");
+    });
+  });
+
+  describe("get_subagent_result", () => {
+    it("reports the cost as its own labelled field", async () => {
+      const { tools } = boot({ showCost: true });
+      runSpending(COST);
+      await spawn(tools);
+      await flush();
+
+      // The agent above ran in the foreground; look it up by the handle its
+      // type gets, which is how the orchestrator would reach it.
+      const text = textOf(await tools.get("get_subagent_result").execute(
+        "tc-2", { agent_id: "general-purpose" }, undefined, undefined, ctx(),
+      ));
+
+      // Pipe-separated `Label: value` fields, matching its neighbours.
+      expect(text).toContain("Cost: ~$0.0123");
+    });
+
+    it("omits the field entirely when unpriced", async () => {
+      const { tools } = boot({ showCost: true });
+      runSpending(0);
+      await spawn(tools);
+      await flush();
+
+      const text = textOf(await tools.get("get_subagent_result").execute(
+        "tc-2", { agent_id: "general-purpose" }, undefined, undefined, ctx(),
+      ));
+
+      expect(text).not.toContain("Cost:");
+    });
+  });
+
+  describe("the background completion notification the model reads", () => {
+    /**
+     * The <task-notification> text sent into the parent conversation. Held
+     * behind a 200ms nudge debounce (plus batch finalization), so this polls
+     * rather than waiting a fixed beat.
+     */
+    async function notificationText(pi: any): Promise<string> {
+      for (let i = 0; i < 60 && pi.sendMessage.mock.calls.length === 0; i++) {
+        await new Promise((r) => setTimeout(r, 25));
+      }
+      return pi.sendMessage.mock.calls.at(-1)?.[0]?.content ?? "";
+    }
+
+    const spawnBackground = (tools: Map<string, any>) =>
+      tools.get("Agent").execute(
+        "tc-1",
+        { prompt: "go", description: "spend", subagent_type: "general-purpose", run_in_background: true },
+        undefined, undefined, ctx(),
+      );
+
+    it("includes the cost in the usage block when enabled", async () => {
+      const { pi, tools } = boot({ showCost: true, defaultJoinMode: "async" });
+      runSpending(COST);
+
+      await spawnBackground(tools);
+
+      expect(await notificationText(pi)).toContain("<estimated_cost_usd>0.0123</estimated_cost_usd>");
+    });
+
+    it("omits it when disabled — this is LLM context, not a display", async () => {
+      // A figure the orchestrator was not asked to track is one it may start
+      // reporting unprompted, so the setting gates the context too, not just
+      // what a human sees.
+      const { pi, tools } = boot({ showCost: false, defaultJoinMode: "async" });
+      runSpending(COST);
+
+      await spawnBackground(tools);
+
+      const text = await notificationText(pi);
+      expect(text).toContain("<total_tokens>");
+      expect(text).not.toContain("estimated_cost_usd");
+    });
+
+    it("omits it for a model with no pricing data", async () => {
+      const { pi, tools } = boot({ showCost: true, defaultJoinMode: "async" });
+      runSpending(0);
+
+      await spawnBackground(tools);
+
+      expect(await notificationText(pi)).not.toContain("estimated_cost_usd");
+    });
+  });
+
+  describe("the completion notification", () => {
+    /** Render a notification through the extension's registered renderer. */
+    function render(pi: any, details: any): string {
+      const [, renderer] = pi.registerMessageRenderer.mock.calls[0];
+      const theme = { fg: (_c: string, s: string) => s, bold: (s: string) => s };
+      return renderer({ details }, { expanded: false }, theme).render().join("\n");
+    }
+
+    const agent = (description: string, totalTokens: number, totalCost: number) => ({
+      id: description, description, status: "completed", toolUses: 1, turnCount: 1,
+      totalTokens, totalCost, durationMs: 1000, resultPreview: "done",
+    });
+
+    it("totals a group, so nobody adds four figures by hand", () => {
+      const { pi } = boot({ showCost: true });
+      const out = render(pi, { ...agent("first", 1000, 0.01), others: [agent("second", 3000, 0.02)] });
+
+      expect(out).toContain("2 agents · 4.0k token · ~$0.03");
+    });
+
+    it("does not total a single agent — the line above already says it", () => {
+      const { pi } = boot({ showCost: true });
+      const out = render(pi, agent("only", 1000, 0.01));
+
+      expect(out).toContain("~$0.01");
+      expect(out).not.toContain("1 agents");
+    });
+
+    it("shows no total, and no per-agent cost, when unpriced", () => {
+      const { pi } = boot({ showCost: true });
+      const out = render(pi, { ...agent("first", 1000, 0), others: [agent("second", 3000, 0)] });
+
+      expect(out).not.toContain("$");
+    });
+
+    it("shows nothing when the setting is off", () => {
+      const { pi } = boot({ showCost: false });
+      const out = render(pi, { ...agent("first", 1000, 0.01), others: [agent("second", 3000, 0.02)] });
+
+      expect(out).not.toContain("$");
+    });
+  });
+});