my-pi/pi-notify/tests/index.test.ts

import assert from "node:assert/strict";
import { afterEach, describe, it } from "node:test";
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { createPiNotifyExtension } from "../src/index.ts";
import type { NotifyConfig } from "../src/types.ts";

type Handler = (...args: any[]) => any;

const savedEnv = new Map(
  ["PI_NOTIFY_MIN_SECONDS", "PI_NOTIFY_MESSAGE_SOURCE"].map((key) => [key, process.env[key]]),
);

afterEach(() => {
  for (const [key, value] of savedEnv) {
    if (value === undefined) delete process.env[key];
    else process.env[key] = value;
  }
});

function harness(mode: ExtensionContext["mode"] = "tui") {
  const handlers = new Map<string, Handler>();
  const eventHandlers = new Map<string, Handler>();
  const commands = new Map<string, { handler: Handler }>();
  const notifications: NotifyConfig[] = [];
  const uiMessages: string[] = [];
  let now = 1_000;

  const pi = {
    on(name: string, handler: Handler) {
      handlers.set(name, handler);
    },
    events: {
      on(name: string, handler: Handler) {
        eventHandlers.set(name, handler);
        return () => eventHandlers.delete(name);
      },
    },
    registerCommand(name: string, command: { handler: Handler }) {
      commands.set(name, command);
    },
  } as unknown as ExtensionAPI;

  createPiNotifyExtension({
    now: () => now,
    createNotificationId: () => "pi-session-test",
    notify: (config) => notifications.push(config),
    getBackend: () => "kitty",
  })(pi);

  const ctx = {
    mode,
    hasUI: mode === "tui" || mode === "rpc",
    cwd: "/tmp/project-a",
    model: { provider: "openai", id: "gpt-test" },
    ui: {
      notify(message: string) {
        uiMessages.push(message);
      },
    },
  } as unknown as ExtensionContext;

  return {
    handlers,
    eventHandlers,
    commands,
    notifications,
    uiMessages,
    ctx,
    setNow(value: number) {
      now = value;
    },
  };
}

describe("pi-notify extension", () => {
  it("notifies once only after retries and follow-ups have settled", async () => {
    const h = harness();
    await h.handlers.get("agent_start")?.({}, h.ctx);
    h.setNow(2_000);
    await h.handlers.get("agent_end")?.({ messages: [{ role: "assistant", content: "retrying" }] }, h.ctx);
    await h.handlers.get("agent_start")?.({}, h.ctx);
    await h.handlers.get("tool_execution_end")?.({ isError: true }, h.ctx);
    h.setNow(5_000);
    await h.handlers.get("agent_end")?.({ messages: [{ role: "assistant", content: "finally done" }] }, h.ctx);

    assert.equal(h.notifications.length, 0, "agent_end must not notify before Pi is settled");
    await h.handlers.get("agent_settled")?.({}, h.ctx);

    assert.equal(h.notifications.length, 1);
    assert.equal(h.notifications[0]?.body, "finally done");
    assert.equal(h.notifications[0]?.durationMs, 4_000, "duration includes retries");
    assert.equal(h.notifications[0]?.notificationId, "pi-session-test");
  });

  it("does not write terminal notifications in RPC or non-interactive modes", async () => {
    process.env.PI_NOTIFY_MIN_SECONDS = "0";
    for (const mode of ["rpc", "json", "print"] as const) {
      const h = harness(mode);
      await h.handlers.get("session_start")?.({}, h.ctx);
      await h.handlers.get("agent_start")?.({}, h.ctx);
      await h.handlers.get("agent_settled")?.({}, h.ctx);
      h.eventHandlers.get("permissions:ui_prompt")?.({
        requestId: "perm-non-tui",
        surface: "bash",
        value: "git push",
      });
      assert.equal(h.notifications.length, 0, mode);
    }
  });

  it("registers session controls and a test notification", async () => {
    const h = harness();
    const command = h.commands.get("notify");
    assert.ok(command);

    await command.handler("off", h.ctx);
    assert.match(h.uiMessages.at(-1) ?? "", /off/);
    await command.handler("test", h.ctx);
    assert.equal(h.notifications.length, 1);
    assert.match(h.notifications[0]?.body ?? "", /click to focus this window/);
    await command.handler("status", h.ctx);
    assert.match(h.uiMessages.at(-1) ?? "", /kitty/);
  });

  it("notifies immediately when permission confirmation needs human input", async () => {
    process.env.PI_NOTIFY_MIN_SECONDS = "3600";
    const h = harness();
    await h.handlers.get("session_start")?.({}, h.ctx);

    h.eventHandlers.get("permissions:ui_prompt")?.({
      requestId: "perm-1",
      source: "tool_call",
      surface: "bash",
      value: "git   push\norigin main",
      agentName: null,
      forwarding: { requesterAgentName: "reviewer", requesterSessionId: "child-1" },
    });

    assert.equal(h.notifications.length, 1);
    assert.equal(h.notifications[0]?.body, "Permission required · reviewer · bash\ngit push origin main");
    assert.equal(h.notifications[0]?.durationMs, 0);
    assert.equal(h.notifications[0]?.minDurationMs, 0, "permission prompts bypass the completion threshold");
    assert.equal(h.notifications[0]?.notificationId, "pi-session-test");
  });

  it("keeps permission details private and honors the session off override", async () => {
    process.env.PI_NOTIFY_MESSAGE_SOURCE = "none";
    const h = harness();
    await h.handlers.get("session_start")?.({}, h.ctx);
    const prompt = { requestId: "perm-secret", surface: "bash", value: "cat private-file" };

    h.eventHandlers.get("permissions:ui_prompt")?.(prompt);
    assert.equal(h.notifications[0]?.body, "Permission confirmation required");

    await h.commands.get("notify")?.handler("off", h.ctx);
    h.eventHandlers.get("permissions:ui_prompt")?.(prompt);
    assert.equal(h.notifications.length, 1);
  });

  it("ignores malformed permission broadcasts and unsubscribes on shutdown", async () => {
    const h = harness();
    await h.handlers.get("session_start")?.({}, h.ctx);
    h.eventHandlers.get("permissions:ui_prompt")?.({ surface: "bash", value: "git push" });
    assert.equal(h.notifications.length, 0);

    await h.handlers.get("session_shutdown")?.({}, h.ctx);
    assert.equal(h.eventHandlers.has("permissions:ui_prompt"), false);
  });
});