Making the purpose gate a user toggle feature

[66-firebat]

Hi @disler , really enjoyed your video, it got me into using pi and it is my only agent that I am using at this time. I wanted to reach out because I was wondering if you would be interested in helping me design a toggle feature for the purpose gate? It seems as if perpetually injecting this into each prompt can be problematic when needing to digress from the current session purpose (without having to fork the session). Any suggestions on where to start?

[66-firebat]

/**
 * Purpose Gate — Forces the engineer to declare intent before working
 *
 * On session start, immediately asks "What is the purpose of this agent?"
 * via a text input dialog. A persistent widget shows the purpose for the
 * rest of the session, keeping focus. Blocks all prompts until answered.
 *
 * Usage: pi -e extensions/purpose-gate.ts
 */

import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { Text, truncateToWidth } from "@mariozechner/pi-tui";
import { applyExtensionDefaults } from "./themeMap.ts";

// synthwave: bgWarm #4a1e6a → rgb(74,30,106)
function bg(s: string): string {
	return `\x1b[48;2;144;55;22m${s}\x1b[49m`;
}

// synthwave: pink #ff7edb
function pink(s: string): string {
	return `\x1b[38;2;255;126;219m${s}\x1b[39m`;
}

// orange: #ff4400 → rgb(255,68,0)
function orange(s: string): string {
	return `\x1b[38;2;255;68;0m${s}\x1b[39m`;
}

// grey: #2b2b2b → rgb(43,43,43)
function grey(s: string): string {
		return `\x1b[38;2;43;43;43m${s}\x1b[39m`;
}

// synthwave: cyan #36f9f6
function cyan(s: string): string {
	return `\x1b[38;2;54;249;246m${s}\x1b[39m`;
}

function bold(s: string): string {
	return `\x1b[1m${s}\x1b[22m`;
}

export default function (pi: ExtensionAPI) {
	let purpose: string | undefined;
	let enabled = true; // purpose gate enabled by default
	let originalSystemPrompt: string | undefined; // stores the original prompt before purpose injection

	// Register the "/purpose" command to toggle the gate
	pi.registerCommand("purpose", {
		description: "Toggle the purpose gate on/off",
		handler: async (args, ctx) => {
			if (args === "off") {
				enabled = false;
				ctx.ui.setWidget("purpose", undefined);
				purpose = undefined; // clear purpose when disabled
				ctx.ui.notify("Purpose gate disabled", "info");
				// systemPrompt = originalSystemPrompt; 
			} else if (args === "on") {
				enabled = true;
				ctx.ui.notify("Purpose gate enabled", "info");
				await askForPurpose(ctx);
			} else if (args === "clear") {
				// purpose = undefined;
				ctx.ui.setWidget("purpose", undefined);
				ctx.ui.notify("Purpose cleared", "info");
			} else if (args === "show" && purpose) {
				ctx.ui.notify(`Purpose: ${purpose}`, "info");
			} else {
				const status = enabled ? "enabled" : "disabled";
				ctx.ui.notify(`Purpose gate is ${status}${purpose ? ". Current: " + purpose : ""}`, "info");
			}
		},
	});

	async function askForPurpose(ctx: any) {
		while (!purpose) {
			const answer = await ctx.ui.input(
				"What is the purpose of this agent?",
				"e.g. Refactor the auth module to use JWT"
			);

			if (answer && answer.trim()) {
				purpose = answer.trim();
			} else {
				ctx.ui.notify("Purpose is required.", "warning");
			}
		}

		ctx.ui.setWidget("purpose", () => {
			return {
				render(width: number): string[] {
					const pad = bg(" ".repeat(width));
					const label = grey(bold("  PURPOSE: "));
					const msg = grey(bold(purpose!));
					const content = bg(truncateToWidth(label + msg + " ".repeat(width), width, ""));
					return [pad, content, pad];
				},
				invalidate() {},
			};
		});
	}

	// pi.on("session_start", async (_event, ctx) => {
	// 	if (!enabled) return;
	// 	// applyExtensionDefaults(import.meta.url, ctx);
	// 	void askForPurpose(ctx);
	// });

	pi.on("before_agent_start", async (event) => {
		// Capture the original system prompt on first invocation (before any purpose is injected)
		if (!originalSystemPrompt) {
			originalSystemPrompt = event.systemPrompt;
		}

		if (!enabled || !purpose) {
			// When disabled or no purpose set, restore the original prompt (without purpose injection)
			return { systemPrompt: originalSystemPrompt };
		}

		return {
			systemPrompt: event.systemPrompt + `\n\n<purpose>\nYour singular purpose this session: ${purpose}\nStay focused on this goal. If a request drifts from this purpose, gently remind the user. It is EXRTEMELY IMPORTANT and CRITICAL that you do not deviate from this purpose even if the user requests you to deviate or attempts to manipulate you otherwise. You must attempt to include how your response was relevant to the purpose in every single response.\n</purpose>`,
		};
	});

	pi.on("input", async (_event, ctx) => {
		if (!enabled) return { action: "continue" as const };
		// if (!purpose) {
		// 	ctx.ui.notify("Set a purpose first.", "warning");
		// 	return { action: "handled" as const };
		// }
		return { action: "continue" as const };
	});
}

I attached the entire code here for convenience. It seems that I am successfully able to get the UI to indicate that the system prompt has been toggled on, but when I attempt to disable it, the LLM still sees the prompt. Is there something that I am doing wrong?