feat: add pi extension (#1)

pi-extension/test/extension.test.js

@@ -0,0 +1,119 @@
+import assert from "node:assert/strict";
+import { mkdtempSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import test from "node:test";
+
+import ponytailExtension from "../index.js";
+
+function createPiHarness() {
+  const events = new Map();
+  const commands = new Map();
+  const appendedEntries = [];
+  const sentUserMessages = [];
+
+  const pi = {
+    on(eventName, handler) {
+      events.set(eventName, handler);
+    },
+    registerCommand(name, options) {
+      commands.set(name, options);
+    },
+    appendEntry(customType, data) {
+      appendedEntries.push({ customType, data });
+    },
+    sendUserMessage(text, options) {
+      sentUserMessages.push({ text, options });
+    },
+  };
+
+  ponytailExtension(pi);
+  return { events, commands, appendedEntries, sentUserMessages };
+}
+
+function createCommandContext(overrides = {}) {
+  return {
+    isIdle: () => true,
+    sessionManager: { getEntries: () => [] },
+    ui: { notify() {} },
+    ...overrides,
+  };
+}
+
+function withTempConfig(fn) {
+  const tempConfigHome = mkdtempSync(join(tmpdir(), "ponytail-test-"));
+  const previousXdg = process.env.XDG_CONFIG_HOME;
+  process.env.XDG_CONFIG_HOME = tempConfigHome;
+
+  return Promise.resolve()
+    .then(fn)
+    .finally(() => {
+      if (previousXdg === undefined) delete process.env.XDG_CONFIG_HOME;
+      else process.env.XDG_CONFIG_HOME = previousXdg;
+      rmSync(tempConfigHome, { recursive: true, force: true });
+    });
+}
+
+test("extension registers Ponytail commands", () => {
+  const { commands } = createPiHarness();
+
+  assert.deepEqual([...commands.keys()].sort(), ["ponytail", "ponytail-help", "ponytail-review"]);
+});
+
+test("/ponytail updates session mode and injects instructions", async () => withTempConfig(async () => {
+  const { commands, events, appendedEntries } = createPiHarness();
+  const ctx = createCommandContext();
+
+  await events.get("session_start")({ reason: "startup" }, ctx);
+  await commands.get("ponytail").handler("ultra", ctx);
+
+  assert.deepEqual(appendedEntries.at(-1), {
+    customType: "ponytail-mode",
+    data: { mode: "ultra" },
+  });
+
+  const result = await events.get("before_agent_start")({ systemPrompt: "BASE" }, ctx);
+  assert.ok(result.systemPrompt.includes("PONYTAIL MODE ACTIVE"));
+  assert.ok(result.systemPrompt.includes("ultra"));
+}));
+
+test("session_start restores latest persisted mode", async () => withTempConfig(async () => {
+  const { events } = createPiHarness();
+  const ctx = createCommandContext({
+    sessionManager: {
+      getEntries: () => [
+        { type: "custom", customType: "ponytail-mode", data: { mode: "lite" } },
+      ],
+    },
+  });
+
+  await events.get("session_start")({ reason: "resume" }, ctx);
+  const result = await events.get("before_agent_start")({ systemPrompt: "BASE" }, ctx);
+
+  assert.ok(result.systemPrompt.includes("lite"));
+}));
+
+test("skill alias commands delegate to Pi skill commands", async () => {
+  const { commands, sentUserMessages } = createPiHarness();
+  const ctx = createCommandContext();
+
+  await commands.get("ponytail-review").handler("", ctx);
+  await commands.get("ponytail-help").handler("", ctx);
+
+  assert.deepEqual(sentUserMessages.map((entry) => entry.text), [
+    "/skill:ponytail-review",
+    "/skill:ponytail-help",
+  ]);
+});
+
+test("normal mode disables persistent instructions", async () => withTempConfig(async () => {
+  const { commands, events } = createPiHarness();
+  const ctx = createCommandContext();
+
+  await events.get("session_start")({ reason: "startup" }, ctx);
+  await commands.get("ponytail").handler("ultra", ctx);
+  await events.get("input")({ text: "normal mode", source: "interactive" }, ctx);
+
+  const disabled = await events.get("before_agent_start")({ systemPrompt: "BASE" }, ctx);
+  assert.equal(disabled, undefined);
+}));