pi-config/extensions/caveman/index.ts

/**
 * pi-caveman — Compressed-output mode for Pi coding agent.
 *
 * Registers a `/caveman` slash command that toggles caveman-speak
 * compression levels and persists state across session restarts.
 *
 * Usage:
 *   /caveman             — toggle between off and full
 *   /caveman lite|full|ultra|wenyan-lite|wenyan-full|wenyan-ultra|off|status
 */

import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import * as fs from "node:fs";
import * as path from "node:path";

// ── Types ────────────────────────────────────────────────────────────────

type CavemanLevel = "lite" | "full" | "ultra" | "wenyan-lite" | "wenyan-full" | "wenyan-ultra";

interface CavemanState {
  active: boolean;
  level: CavemanLevel;
}

// ── Constants ────────────────────────────────────────────────────────────

const DEFAULT_LEVEL: CavemanLevel = "full";

const STATE_DIR = path.join(
  process.env.XDG_STATE_HOME ?? path.join(process.env.HOME!, ".local", "state"),
  "pi-caveman",
);

const STATE_FILE = path.join(STATE_DIR, "caveman-state.json");

const CAVEMAN_SKILL_PATH = path.join(
  process.env.HOME!,
  "ai-assets",
  "skills",
  "caveman",
  "SKILL.md",
);

// ── State helpers ────────────────────────────────────────────────────────

function readState(): CavemanState {
  try {
    const raw = fs.readFileSync(STATE_FILE, "utf-8");
    const parsed = JSON.parse(raw);
    return {
      active: typeof parsed.active === "boolean" ? parsed.active : false,
      level: isValidLevel(parsed.level) ? parsed.level : DEFAULT_LEVEL,
    };
  } catch {
    return { active: false, level: DEFAULT_LEVEL };
  }
}

function writeState(state: CavemanState): void {
  fs.mkdirSync(STATE_DIR, { recursive: true });
  fs.writeFileSync(STATE_FILE, JSON.stringify(state, null, 2), "utf-8");
}

function isValidLevel(v: unknown): v is CavemanLevel {
  return (
    typeof v === "string" &&
    ["lite", "full", "ultra", "wenyan-lite", "wenyan-full", "wenyan-ultra"].includes(v)
  );
}

// ── Skill content cache ──────────────────────────────────────────────────

let cachedSkillContent: string | null = null;

function getSkillContent(): string {
  if (cachedSkillContent !== null) return cachedSkillContent;
  try {
    cachedSkillContent = fs.readFileSync(CAVEMAN_SKILL_PATH, "utf-8");
    return cachedSkillContent;
  } catch {
    return "";
  }
}

// ── Extension ────────────────────────────────────────────────────────────

export default function (pi: ExtensionAPI): void {
  // ── State ──
  let state = readState();

  // ── Footer badge ──
  function updateFooter(ctx: { ui: { setStatus: (k: string, v: string) => void } }): void {
    ctx.ui.setStatus("caveman", state.active ? `🦴 caveman:${state.level}` : "");
  }

  // ── /caveman command ────────────────────────────────────────────────

  pi.registerCommand("caveman", {
    description: "Toggle caveman compressed-output mode (lite|full|ultra|wenyan-lite|wenyan-full|wenyan-ultra|off|status)",
    handler: async (args, ctx) => {
      const arg = args?.trim().toLowerCase() ?? "";

      switch (arg) {
        case "":
          // Toggle: off → full, anything → off
          if (state.active) {
            state.active = false;
            ctx.ui.notify("🦴 Caveman: off", "info");
          } else {
            state.active = true;
            state.level = DEFAULT_LEVEL;
            ctx.ui.notify(`🦴 Caveman: ${state.level}`, "info");
          }
          break;

        case "off":
        case "disable":
          state.active = false;
          ctx.ui.notify("🦴 Caveman: off", "info");
          break;

        case "status":
          ctx.ui.notify(
            state.active ? `🦴 Caveman: ${state.level} (active)` : "🦴 Caveman: off",
            "info",
          );
          return; // Don't update footer (already correct)

        case "lite":
        case "full":
        case "ultra":
        case "wenyan-lite":
        case "wenyan-full":
        case "wenyan-ultra":
          state.active = true;
          state.level = arg as CavemanLevel;
          ctx.ui.notify(`🦴 Caveman: ${state.level}`, "info");
          break;

        default:
          ctx.ui.notify(
            `🦴 Unknown level: "${arg}". Use: lite, full, ultra, wenyan-lite, wenyan-full, wenyan-ultra, off, status`,
            "error",
          );
          return;
      }

      writeState(state);
      updateFooter(ctx);

      // Inject skill content as a user message when activating
      if (state.active) {
        const skill = getSkillContent();
        if (skill) {
          pi.sendUserMessage(
            `/skill:caveman`,
            { deliverAs: "nextTurn" },
          );
        }
      }
    },
  });

  // ── Startup ─────────────────────────────────────────────────────────

  pi.on("session_start", async (_event, ctx) => {
    // Re-read state (may have changed externally)
    state = readState();
    updateFooter(ctx);

    // If active, inject skill into agent context
    if (state.active) {
      const skill = getSkillContent();
      if (skill) {
        pi.appendEntry("caveman-active", { level: state.level });
      }
    }
  });

  // ── Inject caveman rules when active ────────────────────────────────

  pi.on("before_agent_start", async (_event) => {
    // Re-read state in case it changed
    const current = readState();
    if (!current.active) return;

    const skill = getSkillContent();
    if (!skill) return;

    return {
      systemPrompt: `${_event.systemPrompt}\n\n---\nCRITICAL OUTPUT MODE: You are in CAVEMAN MODE at level "${current.level}". Follow these rules exactly:\n\n${skill}`,
    };
  });

  // ── Cleanup on shutdown ─────────────────────────────────────────────

  pi.on("session_shutdown", async () => {
    writeState(state);
  });

}