1// SPDX-FileCopyrightText: Amolith <amolith@secluded.site>
2//
3// SPDX-License-Identifier: Unlicense
4
5/**
6 * Pi Personas Extension
7 *
8 * Switch between agent personas stored as markdown files.
9 *
10 * Personas live in $PI_CODING_AGENT_DIR/personas/ as .md files.
11 * The active persona is appended to the system prompt in a <persona>
12 * section, independent of SYSTEM.md and AGENTS.md.
13 *
14 * Commands:
15 * /persona — Interactive picker to switch persona
16 * /persona none|unset|clear — Remove active persona
17 *
18 * The active persona persists across sessions in a plain text file
19 * at $PI_CODING_AGENT_DIR/persona.
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 getPersonasDir(): string {
40 return path.join(getAgentDir(), "personas");
41}
42
43function getActivePersonaPath(): string {
44 return path.join(getAgentDir(), "persona");
45}
46
47/** List available persona names (without .md extension). */
48function listPersonas(): string[] {
49 const dir = getPersonasDir();
50 if (!fs.existsSync(dir)) return [];
51
52 try {
53 return fs
54 .readdirSync(dir)
55 .filter((f) => f.endsWith(".md"))
56 .map((f) => f.slice(0, -3))
57 .sort();
58 } catch (err) {
59 console.error(`Failed to list personas: ${err}`);
60 return [];
61 }
62}
63
64/**
65 * Read the persisted active persona name, or null if none.
66 * Validates the name against available personas to prevent
67 * path traversal via a tampered persistence file.
68 */
69function readActivePersona(): string | null {
70 const filePath = getActivePersonaPath();
71 if (!fs.existsSync(filePath)) return null;
72
73 try {
74 const name = fs.readFileSync(filePath, "utf-8").trim();
75 if (!name) return null;
76
77 // Reject names that could escape the personas directory
78 const available = listPersonas();
79 if (!available.includes(name)) {
80 console.error(`Persisted persona "${name}" not in available list, ignoring`);
81 return null;
82 }
83
84 return name;
85 } catch (err) {
86 console.error(`Failed to read active persona: ${err}`);
87 return null;
88 }
89}
90
91/** Persist the active persona name (or remove the file to clear). */
92function writeActivePersona(name: string | null): void {
93 const filePath = getActivePersonaPath();
94 try {
95 if (name) {
96 fs.writeFileSync(filePath, `${name}\n`, "utf-8");
97 } else if (fs.existsSync(filePath)) {
98 fs.unlinkSync(filePath);
99 }
100 } catch (err) {
101 console.error(`Failed to write active persona: ${err}`);
102 }
103}
104
105/** Read persona markdown content, or null if file missing. */
106function readPersonaContent(name: string): string | null {
107 const filePath = path.join(getPersonasDir(), `${name}.md`);
108 if (!fs.existsSync(filePath)) return null;
109 try {
110 return fs.readFileSync(filePath, "utf-8");
111 } catch (err) {
112 console.error(`Failed to read persona "${name}": ${err}`);
113 return null;
114 }
115}
116
117const PERSONA_TAG_OPEN =
118 '<persona section_description="This section defines your personality. Fully embody the character, voice, and behaviour described within.">';
119const PERSONA_TAG_CLOSE = "</persona>";
120
121function wrapPersona(content: string): string {
122 return `\n\n${PERSONA_TAG_OPEN}\n${content}\n${PERSONA_TAG_CLOSE}\n`;
123}
124
125export default function (pi: ExtensionAPI) {
126 // "active" tracks what the user last picked (persisted to disk).
127 // "session" is locked at session_start and never changes — it
128 // controls what goes into the system prompt for this session.
129 let activePersona: string | null = null;
130 let sessionPersonaContent: string | null = null;
131 let pendingReminder: string | null = null;
132
133 /** Set the active persona name. Does NOT touch session prompt. */
134 function setActive(name: string | null): void {
135 activePersona = name;
136 writeActivePersona(name);
137 }
138
139 // Lock persona into the system prompt at session start.
140 // This content never changes for the lifetime of the session.
141 pi.on("session_start", async (_event, ctx) => {
142 const persisted = readActivePersona();
143 pendingReminder = null;
144 sessionPersonaContent = null;
145
146 if (persisted) {
147 const content = readPersonaContent(persisted);
148 if (content) {
149 activePersona = persisted;
150 sessionPersonaContent = content;
151 ctx.ui.notify(`Set persona to "${persisted}"`, "info");
152 } else {
153 // Persona file was removed since last session
154 setActive(null);
155 }
156 }
157 });
158
159 // Append the session persona to the system prompt (ephemeral,
160 // never persisted). The content is identical every turn so the
161 // provider cache stays valid.
162 //
163 // Mid-session switches only fire a one-shot message — they never
164 // alter the system prompt.
165 pi.on("before_agent_start", async (event) => {
166 if (!sessionPersonaContent && !pendingReminder) return undefined;
167
168 const result: {
169 systemPrompt?: string;
170 message?: { customType: string; content: string; display: boolean };
171 } = {};
172
173 if (sessionPersonaContent) {
174 result.systemPrompt = event.systemPrompt + wrapPersona(sessionPersonaContent);
175 }
176
177 if (pendingReminder) {
178 result.message = {
179 customType: "persona-switch",
180 content: pendingReminder,
181 display: false,
182 };
183 pendingReminder = null;
184 }
185
186 return result;
187 });
188
189 // Register /persona command
190 pi.registerCommand("persona", {
191 description: "Switch agent persona or clear it",
192 getArgumentCompletions: (prefix: string) => {
193 const names = [...listPersonas(), ...CLEAR_KEYWORDS];
194 const filtered = names.filter((n) => n.startsWith(prefix)).map((n) => ({ value: n, label: n }));
195 return filtered.length > 0 ? filtered : null;
196 },
197 handler: async (args, ctx) => {
198 const arg = args?.trim() || "";
199
200 // /persona none|unset|clear — remove active persona
201 if (CLEAR_KEYWORDS.includes(arg.toLowerCase())) {
202 if (!activePersona) {
203 ctx.ui.notify("No persona is active", "info");
204 return;
205 }
206 const was = activePersona;
207 pendingReminder = `Persona "${was}" has been cleared. Revert to your default behaviour.`;
208 setActive(null);
209 ctx.ui.notify(`Persona cleared (was: ${was})`, "info");
210 return;
211 }
212
213 // /persona someName — direct switch (not in picker)
214 if (arg) {
215 const available = listPersonas();
216 if (!available.includes(arg)) {
217 ctx.ui.notify(`Persona "${arg}" not found`, "error");
218 return;
219 }
220 const content = readPersonaContent(arg);
221 if (!content) {
222 ctx.ui.notify(`Persona "${arg}" could not be read`, "error");
223 return;
224 }
225 pendingReminder = wrapPersona(content);
226 setActive(arg);
227 ctx.ui.notify(`Set persona to "${arg}"`, "info");
228 return;
229 }
230
231 // /persona — interactive picker
232 const available = listPersonas();
233 if (available.length === 0) {
234 ctx.ui.notify("No personas available", "info");
235 return;
236 }
237
238 const options = available.map((name) => (name === activePersona ? `${name} (active)` : name));
239 const unsetLabel = activePersona ? "unset" : "unset (active)";
240 options.unshift(unsetLabel);
241
242 const choice = await ctx.ui.select("Pick a persona:", options);
243
244 if (!choice) return; // cancelled
245
246 // Strip " (active)" suffix to get the real name
247 const picked = choice.replace(/ \(active\)$/, "");
248
249 if (picked === "unset") {
250 if (!activePersona) {
251 ctx.ui.notify("No persona is active", "info");
252 } else {
253 const was = activePersona;
254 pendingReminder = `Persona "${was}" has been cleared. Revert to your default behaviour.`;
255 setActive(null);
256 ctx.ui.notify(`Persona cleared (was: ${was})`, "info");
257 }
258 return;
259 }
260
261 const content = readPersonaContent(picked);
262 if (!content) {
263 ctx.ui.notify(`Persona "${picked}" could not be read`, "error");
264 return;
265 }
266 pendingReminder = wrapPersona(content);
267 setActive(picked);
268 ctx.ui.notify(`Set persona to "${picked}"`, "info");
269 },
270 });
271}