my-pi/pi-subagents/test/bundled-agents.test.ts

/**
 * bundled-agents.test.ts — `bundledAgentsDir` extension option: bundle-shipped
 * agent .md files register as embedded defaults (isDefault, no source/path),
 * sit between embedded defaults and user agents in precedence, and stay out of
 * the registry when the option is unset.
 */
import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { buildAgentRegistry, getAgentConfig, getDefaultAgentNames, registerAgents } from "../src/agent-types.js";
import { loadBundledAgents } from "../src/custom-agents.js";
import { createPiSubagentsExtension } from "../src/index.js";

let tmpDir: string;
let bundledDir: string;
let originalHome: string | undefined;
let originalAgentDir: string | undefined;

beforeEach(() => {
  tmpDir = mkdtempSync(join(tmpdir(), "pi-bundled-agents-"));
  bundledDir = join(tmpDir, "bundled");
  mkdirSync(bundledDir, { recursive: true });
  originalHome = process.env.HOME;
  originalAgentDir = process.env.PI_CODING_AGENT_DIR;
  process.env.HOME = join(tmpDir, "home");
  delete process.env.PI_CODING_AGENT_DIR;
});

afterEach(() => {
  if (originalHome == null) delete process.env.HOME;
  else process.env.HOME = originalHome;
  if (originalAgentDir == null) delete process.env.PI_CODING_AGENT_DIR;
  else process.env.PI_CODING_AGENT_DIR = originalAgentDir;
  rmSync(tmpDir, { recursive: true, force: true });
  registerAgents(new Map());
});

function writeBundled(name: string, content: string) {
  writeFileSync(join(bundledDir, `${name}.md`), content);
}

describe("loadBundledAgents", () => {
  it("parses agent files with the custom-agent parser and stamps them as built-ins", () => {
    writeBundled("scout", `---
name: scout
description: Bundle scout
prompt_mode: append
tools: read, grep
---

Bundle prompt body.`);

    const agents = loadBundledAgents(bundledDir);
    const scout = agents.get("scout");
    expect(scout).toBeDefined();
    expect(scout?.description).toBe("Bundle scout");
    expect(scout?.systemPrompt).toBe("Bundle prompt body.");
    expect(scout?.promptMode).toBe("append");
    expect(scout?.builtinToolNames).toEqual(["read", "grep"]);
    expect(scout?.isDefault).toBe(true);
    expect(scout?.source).toBeUndefined();
    expect(scout?.sourcePath).toBeUndefined();
  });

  it("returns an empty map for a missing or empty directory", () => {
    expect(loadBundledAgents(join(tmpDir, "missing")).size).toBe(0);
    expect(loadBundledAgents(bundledDir).size).toBe(0);
  });

  it("skips non-.md files and reserved ':' names", () => {
    writeBundled("fine", "---\ndescription: ok\n---\nbody");
    writeFileSync(join(bundledDir, "notes.txt"), "not an agent");
    writeBundled("bad", "---\nname: 'plugin:scout'\ndescription: bad\n---\nbody");
    const agents = loadBundledAgents(bundledDir);
    expect([...agents.keys()]).toEqual(["fine"]);
  });
});

describe("bundled agent precedence", () => {
  it("replaces a same-name embedded default and reports as a default", () => {
    writeBundled("Explore", `---
name: Explore
description: Bundled Explore override
prompt_mode: replace
---

Bundled Explore body.`);

    const bundled = loadBundledAgents(bundledDir);
    const registry = buildAgentRegistry(new Map(), bundled);
    expect(registry.get("Explore")?.systemPrompt).toBe("Bundled Explore body.");
    expect(registry.get("Explore")?.isDefault).toBe(true);
    // Unrelated embedded defaults are untouched.
    expect(registry.get("general-purpose")?.systemPrompt).not.toBe("Bundled Explore body.");
  });

  it("new bundled names join the default set", () => {
    writeBundled("bundle-only", "---\ndescription: only in bundle\n---\nbody");
    const registry = buildAgentRegistry(new Map(), loadBundledAgents(bundledDir));
    expect(registry.get("bundle-only")?.isDefault).toBe(true);
    registerAgents(new Map(), loadBundledAgents(bundledDir));
    expect(getDefaultAgentNames()).toContain("bundle-only");
    expect(getAgentConfig("bundle-only")?.description).toBe("only in bundle");
  });

  it("user agents still override bundled definitions", () => {
    writeBundled("scout", "---\ndescription: bundled\n---\nbundled body");
    const user = new Map([
      ["scout", {
        name: "scout",
        description: "project override",
        systemPrompt: "project body",
        promptMode: "replace" as const,
        extensions: true as const,
        skills: true as const,
        source: "project" as const,
        sourcePath: "/proj/.pi/agents/scout.md",
      }],
    ]);
    const registry = buildAgentRegistry(user, loadBundledAgents(bundledDir));
    expect(registry.get("scout")?.description).toBe("project override");
    expect(registry.get("scout")?.isDefault).not.toBe(true);
  });
});

describe("extension option validation", () => {
  it("rejects a path that is both mandatory and excluded", () => {
    const dir = mkdtempSync(join(tmpdir(), "pi-subagents-overlap-"));
    const path = join(dir, "permission-system.ts");
    writeFileSync(path, "export default () => {};\n");
    expect(() =>
      createPiSubagentsExtension({
        mandatoryExtensionPaths: [path],
        excludedExtensionPaths: [path],
      }),
    ).toThrow("both mandatory and excluded");
    rmSync(dir, { recursive: true, force: true });
  });

  it("canonicalizes bundledAgentsDir and rejects missing directories", () => {
    expect(() =>
      createPiSubagentsExtension({ bundledAgentsDir: join(tmpDir, "missing") }),
    ).toThrow("is unavailable");
    expect(() => createPiSubagentsExtension({ bundledAgentsDir: bundledDir })).not.toThrow();
  });
});