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 savedMinSeconds = process.env.PI_NOTIFY_MIN_SECONDS;

afterEach(() => {
  if (savedMinSeconds === undefined) delete process.env.PI_NOTIFY_MIN_SECONDS;
  else process.env.PI_NOTIFY_MIN_SECONDS = savedMinSeconds;
});

function harness(mode: ExtensionContext["mode"] = "tui") {
  const handlers = 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);
    },
    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,
    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("agent_start")?.({}, h.ctx);
      await h.handlers.get("agent_settled")?.({}, h.ctx);
      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/);
  });
});