index.ts

  1// SPDX-FileCopyrightText: Amolith <amolith@secluded.site>
  2//
  3// SPDX-License-Identifier: Unlicense
  4
  5/**
  6 * Pi Personalities Extension
  7 *
  8 * Switch between agent personalities stored as markdown files.
  9 *
 10 * Personalities live in $PI_CODING_AGENT_DIR/personalities/ as .md files.
 11 * The active personality is appended to the system prompt in a <personality>
 12 * section, independent of SYSTEM.md and AGENTS.md.
 13 *
 14 * Commands:
 15 *   /personality              — Interactive picker to switch personality
 16 *   /personality none|unset|clear — Remove active personality
 17 *
 18 * The active personality persists across sessions in a plain text file
 19 * at $PI_CODING_AGENT_DIR/personality.
 20 */
 21
 22import * as fs from "node:fs";
 23import * as path from "node:path";
 24import * as os from "node:os";
 25import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
 26
 27const CLEAR_KEYWORDS = ["none", "unset", "clear"];
 28
 29function getAgentDir(): string {
 30	const envDir = process.env.PI_CODING_AGENT_DIR;
 31	if (envDir) {
 32		if (envDir === "~") return os.homedir();
 33		if (envDir.startsWith("~/")) return path.join(os.homedir(), envDir.slice(2));
 34		return envDir;
 35	}
 36	return path.join(os.homedir(), ".pi", "agent");
 37}
 38
 39function getPersonalitiesDir(): string {
 40	return path.join(getAgentDir(), "personalities");
 41}
 42
 43function getActivePersonalityPath(): string {
 44	return path.join(getAgentDir(), "personality");
 45}
 46
 47/** List available personality names (without .md extension). */
 48function listPersonalities(): string[] {
 49	const dir = getPersonalitiesDir();
 50	if (!fs.existsSync(dir)) return [];
 51
 52	return fs
 53		.readdirSync(dir)
 54		.filter((f) => f.endsWith(".md"))
 55		.map((f) => f.slice(0, -3))
 56		.sort();
 57}
 58
 59/** Read the persisted active personality name, or null if none. */
 60function readActivePersonality(): string | null {
 61	const filePath = getActivePersonalityPath();
 62	if (!fs.existsSync(filePath)) return null;
 63
 64	const name = fs.readFileSync(filePath, "utf-8").trim();
 65	return name || null;
 66}
 67
 68/** Persist the active personality name (or remove the file to clear). */
 69function writeActivePersonality(name: string | null): void {
 70	const filePath = getActivePersonalityPath();
 71	if (name) {
 72		fs.writeFileSync(filePath, name + "\n", "utf-8");
 73	} else if (fs.existsSync(filePath)) {
 74		fs.unlinkSync(filePath);
 75	}
 76}
 77
 78/** Read personality markdown content, or null if file missing. */
 79function readPersonalityContent(name: string): string | null {
 80	const filePath = path.join(getPersonalitiesDir(), `${name}.md`);
 81	if (!fs.existsSync(filePath)) return null;
 82	return fs.readFileSync(filePath, "utf-8");
 83}
 84
 85const PERSONALITY_TAG_OPEN =
 86	'<personality section_description="This section defines your personality. Fully embody the character, voice, and behaviour described within.">';
 87const PERSONALITY_TAG_CLOSE = "</personality>";
 88
 89function wrapPersonality(content: string): string {
 90	return `\n\n${PERSONALITY_TAG_OPEN}\n${content}\n${PERSONALITY_TAG_CLOSE}\n`;
 91}
 92
 93export default function (pi: ExtensionAPI) {
 94	let activePersonality: string | null = null;
 95	let personalityContent: string | null = null;
 96	let isFirstTurn = true;
 97
 98	/** Load personality by name. Returns true if found and loaded. */
 99	function loadPersonality(name: string): boolean {
100		const content = readPersonalityContent(name);
101		if (!content) return false;
102
103		activePersonality = name;
104		personalityContent = content;
105		return true;
106	}
107
108	/** Clear active personality (both in-memory and on disk). */
109	function clearPersonality(): void {
110		activePersonality = null;
111		personalityContent = null;
112		writeActivePersonality(null);
113	}
114
115	/** Switch to a personality: load into memory and persist to disk. */
116	function switchPersonality(name: string): boolean {
117		if (!loadPersonality(name)) return false;
118		writeActivePersonality(name);
119		return true;
120	}
121
122	// Restore persisted personality on session start
123	pi.on("session_start", async (_event, ctx) => {
124		const persisted = readActivePersonality();
125
126		// Determine whether this is a fresh session or a resumed one
127		const entries = ctx.sessionManager.getEntries();
128		const hasMessages = entries.some((e) => e.type === "message" && e.message.role === "assistant");
129		isFirstTurn = !hasMessages;
130
131		if (persisted) {
132			if (loadPersonality(persisted)) {
133				ctx.ui.notify(`Set personality to "${persisted}"`, "info");
134			} else {
135				// Personality file was removed since last session
136				clearPersonality();
137			}
138		}
139	});
140
141	// Inject personality into the system prompt or as a reminder message.
142	// On the first turn (no prior messages), we modify the system prompt
143	// directly so the personality is baked into the initial cache.
144	// On subsequent turns, we inject a reminder message instead to avoid
145	// invalidating the inference cache.
146	pi.on("before_agent_start", async (event) => {
147		if (!personalityContent) return undefined;
148
149		if (isFirstTurn) {
150			isFirstTurn = false;
151			return {
152				systemPrompt: event.systemPrompt + wrapPersonality(personalityContent),
153			};
154		}
155
156		return {
157			message: {
158				customType: "system-reminder",
159				content: wrapPersonality(personalityContent),
160				display: false,
161			},
162		};
163	});
164
165	// Register /personality command
166	pi.registerCommand("personality", {
167		description: "Switch agent personality or clear it",
168		getArgumentCompletions: (prefix: string) => {
169			const names = [...listPersonalities(), ...CLEAR_KEYWORDS];
170			const filtered = names.filter((n) => n.startsWith(prefix)).map((n) => ({ value: n, label: n }));
171			return filtered.length > 0 ? filtered : null;
172		},
173		handler: async (args, ctx) => {
174			const arg = args?.trim() || "";
175
176			// /personality none|unset|clear — remove active personality
177			if (CLEAR_KEYWORDS.includes(arg.toLowerCase())) {
178				if (!activePersonality) {
179					ctx.ui.notify("No personality is active", "info");
180					return;
181				}
182				const was = activePersonality;
183				clearPersonality();
184				ctx.ui.notify(`Personality cleared (was: ${was})`, "info");
185				return;
186			}
187
188			// /personality someName — direct switch (not in picker)
189			if (arg) {
190				const available = listPersonalities();
191				if (!available.includes(arg)) {
192					ctx.ui.notify(`Personality "${arg}" not found`, "error");
193					return;
194				}
195				switchPersonality(arg);
196				ctx.ui.notify(`Set personality to "${arg}"`, "info");
197				return;
198			}
199
200			// /personality — interactive picker
201			const available = listPersonalities();
202			if (available.length === 0) {
203				ctx.ui.notify("No personalities available", "info");
204				return;
205			}
206
207			const options = available.map((name) => (name === activePersonality ? `${name} (active)` : name));
208			const unsetLabel = activePersonality ? "unset" : "unset (active)";
209			options.unshift(unsetLabel);
210
211			const choice = await ctx.ui.select("Pick a personality:", options);
212
213			if (!choice) return; // cancelled
214
215			// Strip " (active)" suffix to get the real name
216			const picked = choice.replace(/ \(active\)$/, "");
217
218			if (picked === "unset") {
219				if (!activePersonality) {
220					ctx.ui.notify("No personality is active", "info");
221				} else {
222					const was = activePersonality;
223					clearPersonality();
224					ctx.ui.notify(`Personality cleared (was: ${was})`, "info");
225				}
226				return;
227			}
228
229			switchPersonality(picked);
230			ctx.ui.notify(`Set personality to "${picked}"`, "info");
231		},
232	});
233}