Pi Agent Mentor mode

mentor-mode.ts

/**
 * Mentor Mode Extension
 *
 * Transforms the agent into an expert instructor. While enabled, the agent:
 *  - Guides the user through work as a teacher would
 *  - Asks questions to test understanding as work progresses
 *  - Prompts the user with choices to steer the work
 *  - Corrects obviously incorrect choices and explains why
 *  - Answers "why" questions thoroughly and clearly
 *  - Whenever output could be produced, offers to do it itself or let the user
 *    do it; if the user does it, waits for confirmation, then reads, evaluates,
 *    and continues
 *
 * Commands / controls:
 *   /mentor              Toggle Mentor Mode
 *   Ctrl+Alt+M           Toggle Mentor Mode
 *   --mentor             Start in Mentor Mode
 *
 * Tools exposed to the LLM (only while Mentor Mode is active):
 *   mentor_question      Ask a multiple-choice / understanding question
 *   mentor_offer         Offer to produce output vs. let the user do it
 */

import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";

const MENTOR_TOOLS = ["mentor_question", "mentor_offer"];

const MENTOR_SYSTEM_PROMPT = `

========== MENTOR MODE ACTIVE ==========

You are now operating in MENTOR MODE. Act as an expert instructor guiding the
user through the work at hand. Your goal is not just to complete the task, but
to help the user genuinely understand it.

Adopt these behaviors for the rest of this session:

1. TEACH, DON'T JUST DO. Walk the user through the work step by step as a
   knowledgeable mentor would. Explain your reasoning as you go.

2. TEST UNDERSTANDING. As work progresses, periodically check the user's
   understanding. Use the mentor_question tool to pose focused questions
   (multiple choice is ideal) at natural checkpoints — after explaining a
   concept, after reaching a milestone, or before moving to a new phase.

3. GUIDE WITH CHOICES. When the work could go in more than one direction,
   use mentor_question to present the user with clear options and let them
   steer. Recommend the option you think is best, but respect their choice.

4. CORRECT MISTAKES CLEARLY. When the user selects an option that is obviously
   incorrect, or gives an answer that reveals a misconception, do not just move
   on. Correct them, explain clearly WHY the choice is wrong, and explain why
   the right answer is right. Be encouraging, not condescending.

5. ANSWER "WHY" THOROUGHLY. Whenever the user asks "why" — why something works
   a certain way, why a tool behaves as it does, why an answer is correct —
   give a complete, clear explanation. "Why" questions are the heart of
   learning; never brush them off. Connect the explanation to the broader
   topic so the user builds deeper understanding.

6. OFFER THE KEYBOARD. Whenever an opportunity to produce concrete output
   arises — writing a file, running a command, drafting an answer, solving a
   challenge step — pause FIRST and use the mentor_offer tool to ask whether
   the user wants you to produce it, or whether they want to do it themselves.

   - If they choose for YOU to do it: proceed and produce the output.
   - If they choose to do it THEMSELVES: STOP. Do not produce the output.
     Wait for the user to confirm they have completed it. When they confirm,
     read the relevant output (files they wrote, commands they ran, answers
     they gave), evaluate it, give constructive feedback, and only then
     continue with the work. If their output is incorrect, treat it as a
     teaching moment per rule 4.

7. KEEP THE USER IN THE LOOP. Narrate what you are about to do and why before
   doing it. Avoid long silent runs of tool calls without explanation.

Mentor Mode is a teaching posture. Stay accurate and rigorous — never invent
facts to make a lesson flow. If you are unsure of something, say so and
investigate it together with the user.
========================================
`;

interface MentorQuestionDetails {
	question: string;
	options: string[];
	answer: string | null;
	correctIndex: number | null;
	wasCorrect: boolean | null;
	wasCustom: boolean;
	explanation: string | null;
}

interface MentorOfferDetails {
	task: string;
	choice: "agent" | "user" | null;
}

const QuestionParams = Type.Object({
	question: Type.String({ description: "The question to ask the user." }),
	options: Type.Array(
		Type.Object({
			label: Type.String({ description: "Display label for the option." }),
			description: Type.Optional(
				Type.String({ description: "Optional one-line description shown under the label." }),
			),
		}),
		{ description: "The choices the user can pick from." },
	),
	correctIndex: Type.Optional(
		Type.Integer({
			description:
				"0-based index of the correct option, for understanding-check questions. Omit for open-ended guiding choices where there is no single right answer.",
		}),
	),
	explanation: Type.Optional(
		Type.String({
			description:
				"Why the correct option is correct (and why wrong ones are wrong). Used to correct the user if they pick incorrectly. Provide for understanding-check questions.",
		}),
	),
});

const OfferParams = Type.Object({
	task: Type.String({
		description:
			"A short description of the output that is about to be produced (e.g. 'write the solution script', 'run the nmap scan', 'draft the analysis').",
	}),
});

export default function mentorModeExtension(pi: ExtensionAPI): void {
	let mentorMode = false;

	function updateStatus(ctx: ExtensionContext): void {
		if (mentorMode) {
			ctx.ui.setStatus("mentor-mode", ctx.ui.theme.fg("accent", "🎓 mentor"));
		} else {
			ctx.ui.setStatus("mentor-mode", undefined);
		}
	}

	function persist(): void {
		pi.appendEntry("mentor-mode", { enabled: mentorMode });
	}

	function enableMentorTools(): void {
		const active = pi.getActiveTools();
		pi.setActiveTools([...new Set([...active, ...MENTOR_TOOLS])]);
	}

	function disableMentorTools(): void {
		pi.setActiveTools(pi.getActiveTools().filter((name) => !MENTOR_TOOLS.includes(name)));
	}

	function toggle(ctx: ExtensionContext): void {
		mentorMode = !mentorMode;
		if (mentorMode) {
			enableMentorTools();
			ctx.ui.notify("Mentor Mode enabled — guiding you as an expert instructor.", "info");
		} else {
			disableMentorTools();
			ctx.ui.notify("Mentor Mode disabled.", "info");
		}
		updateStatus(ctx);
		persist();
	}

	pi.registerFlag("mentor", {
		description: "Start in Mentor Mode (expert instructor guidance)",
		type: "boolean",
		default: false,
	});

	pi.registerCommand("mentor", {
		description: "Toggle Mentor Mode (expert instructor guidance)",
		handler: async (_args, ctx) => toggle(ctx),
	});

	// Sanity check: verify ctx.ui.select returns the chosen VALUE (a string),
	// not a numeric index. Run with /mentor-probe after touching dialog code.
	pi.registerCommand("mentor-probe", {
		description: "Probe ctx.ui.select return type (dev sanity check)",
		handler: async (_args, ctx) => {
			if (!ctx.hasUI) {
				ctx.ui.notify("mentor-probe: UI unavailable in this mode.", "warning");
				return;
			}
			const options = ["alpha", "bravo", "charlie"];
			const result = await ctx.ui.select("mentor-probe: pick any option", options);
			if (result === undefined) {
				ctx.ui.notify("mentor-probe: select returned undefined (cancelled). OK.", "info");
				return;
			}
			const isValue = typeof result === "string";
			const isKnownValue = isValue && options.includes(result as string);
			const looksLikeIndex = typeof result === "number";
			const verdict = isKnownValue
				? "PASS: select returned the chosen value."
				: looksLikeIndex
					? `FAIL: select returned a number (${result}) — treated as index. Expected a value string.`
					: `FAIL: select returned ${typeof result} (${JSON.stringify(result)}) — not a known option value.`;
			const level: "info" | "error" = isKnownValue ? "info" : "error";
			ctx.ui.notify(`mentor-probe: ${verdict} [typeof=${typeof result}]`, level);
		},
	});

	pi.registerShortcut("ctrl+alt+m", {
		description: "Toggle Mentor Mode",
		handler: async (ctx) => toggle(ctx),
	});

	// Restore state on session start / resume / reload, and honor --mentor flag.
	pi.on("session_start", async (_event, ctx) => {
		let enabled = false;
		if (pi.getFlag("mentor")) {
			enabled = true;
		} else {
			for (const entry of ctx.sessionManager.getEntries()) {
				if (entry.type === "custom" && entry.customType === "mentor-mode") {
					const data = entry.data as { enabled?: boolean } | undefined;
					if (data?.enabled) {
						enabled = true;
						break;
					}
				}
			}
		}
		mentorMode = enabled;
		if (mentorMode) {
			enableMentorTools();
		}
		updateStatus(ctx);
	});

	// Inject mentor instructions into the system prompt while active.
	pi.on("before_agent_start", async (event) => {
		if (mentorMode) {
			return { systemPrompt: event.systemPrompt + MENTOR_SYSTEM_PROMPT };
		}
		return undefined;
	});

	// --- mentor_question tool ----------------------------------------------
	pi.registerTool({
		name: "mentor_question",
		label: "Mentor Question",
		description:
			"Ask the user a multiple-choice question to test understanding or to let them steer the work. Use it at natural checkpoints while teaching in Mentor Mode.",
		promptSnippet: "Ask the user a multiple-choice question (understanding check or steering choice)",
		promptGuidelines: [
			"Use mentor_question at natural teaching checkpoints to test the user's understanding or to present steering choices. Provide correctIndex and explanation for understanding-check questions so you can correct wrong answers with a clear why.",
		],
		parameters: QuestionParams,
		executionMode: "sequential",

		async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
			const options = Array.isArray(params.options) ? params.options : [];
			const labels = options.map((o) => o.label);
			const correctIndex =
				typeof params.correctIndex === "number" && params.correctIndex >= 0
					? params.correctIndex
					: null;

			if (labels.length === 0) {
				return {
					content: [
						{ type: "text", text: "Error: mentor_question was called with no options." },
					],
					details: {
						question: params.question,
						options: [],
						answer: null,
						correctIndex,
						wasCorrect: null,
						wasCustom: false,
						explanation: params.explanation ?? null,
					} as MentorQuestionDetails,
				};
			}

			// Build display strings with optional descriptions.
			const displayItems = options.map((o) => {
				const desc = o.description ? `  — ${o.description}` : "";
				return `${o.label}${desc}`;
			});
			displayItems.push("Type my own answer…");

			let answer: string | null = null;
			let wasCustom = false;

			if (ctx.hasUI) {
				// ctx.ui.select returns the chosen option STRING (or undefined on
				// cancel), NOT a numeric index. Match it back to the option it came
				// from so labels/descriptions stay in sync.
				const choice = await ctx.ui.select(params.question, displayItems);
				if (choice === undefined) {
					answer = null;
				} else if (choice === displayItems[displayItems.length - 1]) {
					// "Type my own answer…"
					const custom = await ctx.ui.input("Your answer:", "");
					answer = custom ?? null;
					wasCustom = true;
				} else {
					const matchedIdx = displayItems.indexOf(choice);
					answer = matchedIdx >= 0 ? labels[matchedIdx] : choice;
				}
			} else {
				// Non-interactive fallback: report that UI is unavailable.
				return {
					content: [
						{
							type: "text",
							text: "User interaction is unavailable in this mode. Ask the question in plain text instead.",
						},
					],
					details: {
						question: params.question,
						options: labels,
						answer: null,
						correctIndex,
						wasCorrect: null,
						wasCustom: false,
						explanation: params.explanation ?? null,
					} as MentorQuestionDetails,
				};
			}

			const details: MentorQuestionDetails = {
				question: params.question,
				options: labels,
				answer,
				correctIndex,
				wasCorrect: null,
				wasCustom,
				explanation: params.explanation ?? null,
			};

			if (answer === null) {
				return {
					content: [{ type: "text", text: "User cancelled the question." }],
					details,
				};
			}

			// Determine correctness for understanding-check questions.
			if (correctIndex !== null && !wasCustom) {
				const selectedIdx = labels.indexOf(answer);
				details.wasCorrect = selectedIdx === correctIndex;
				const correctLabel = labels[correctIndex];
				if (details.wasCorrect) {
					return {
						content: [
							{
								type: "text",
								text: `User answered: "${answer}" — CORRECT.\n\nAcknowledge briefly and continue.`,
							},
						],
						details,
					};
				}
				return {
					content: [
						{
							type: "text",
							text:
								`User answered: "${answer}" — INCORRECT.\n` +
								`Correct answer: "${correctLabel}".\n` +
								`Explanation: ${params.explanation ?? "(no explanation provided)"}\n\n` +
								`Correct the user, explain clearly why their choice is wrong and why ` +
								`"${correctLabel}" is right, then continue.`,
						},
					],
					details,
				};
			}

			// Open-ended steering choice, or a custom free-text answer.
			if (wasCustom) {
				return {
					content: [
						{
							type: "text",
							text: `User typed their own answer: "${answer}". Evaluate it as an instructor would: if it reveals a misconception, correct it with a clear explanation; otherwise respond and continue.`,
						},
					],
					details,
				};
			}

			return {
				content: [{ type: "text", text: `User selected: "${answer}". Respect their choice and continue.` }],
				details,
			};
		},
	});

	// --- mentor_offer tool -------------------------------------------------
	pi.registerTool({
		name: "mentor_offer",
		label: "Mentor Offer",
		description:
			"Before producing concrete output, ask the user whether they want you to produce it or do it themselves. If they choose to do it themselves, stop and wait for their confirmation, then read and evaluate their work.",
		promptSnippet: "Offer to produce output, or let the user do it themselves",
		promptGuidelines: [
			"Use mentor_offer every time you are about to produce concrete output (write a file, run a command, draft an answer, solve a step). Ask before doing it. If the user chooses to do it themselves, STOP and wait for them to confirm completion; then read their output, evaluate it, give feedback, and continue.",
		],
		parameters: OfferParams,
		executionMode: "sequential",

		async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
			const task = params.task;

			if (!ctx.hasUI) {
				return {
					content: [
						{
							type: "text",
							text: "User interaction is unavailable in this mode. Ask in plain text whether the user wants you to produce the output or do it themselves.",
						},
					],
					details: { task, choice: null } as MentorOfferDetails,
				};
			}

			const selfLabel = "I'll do it myself";
			// ctx.ui.select returns the chosen option STRING (or undefined on
			// cancel), NOT a numeric index.
			const choice = await ctx.ui.select(
				`Mentor Mode — next step:\n${task}\n\nHow would you like to proceed?`,
				[selfLabel, "You do it (Mentor)"],
			);

			const details: MentorOfferDetails = { task, choice: null };

			if (choice === undefined) {
				return {
					content: [{ type: "text", text: "User cancelled the offer." }],
					details,
				};
			}

			if (choice === selfLabel) {
				details.choice = "user";
				return {
					content: [
						{
							type: "text",
							text:
								`The user will do this step themselves: ${task}\n\n` +
								`STOP now. Do NOT produce this output. End your turn and wait for the ` +
								`user to confirm they have finished. When they confirm, read the relevant ` +
								`output (the file(s) they wrote, the command(s) they ran, the answer they ` +
								`gave), evaluate it as an instructor, give constructive feedback, and then ` +
								`continue with the work. If their work is incorrect, treat it as a teaching ` +
								`moment: explain why and how to fix it.`,
						},
					],
					details,
				};
			}

			details.choice = "agent";
			return {
				content: [
					{
						type: "text",
						text: `The user wants you to produce this output: ${task}. Proceed to do it, narrating what you are doing and why.`,
					},
				],
				details,
			};
		},
	});
}