562 lines
17 KiB
TypeScript
562 lines
17 KiB
TypeScript
/**
|
|
* Plan Mode Extension
|
|
*
|
|
* Read-only exploration + structured planning + gated step-by-step execution.
|
|
*
|
|
* Features:
|
|
* - /plan command or Ctrl+Alt+P to toggle read-only plan mode
|
|
* - Bash restricted to allowlisted read-only commands, edit/write disabled
|
|
* - `submit_plan` tool: model returns a STRUCTURED plan (title + steps + risks + files)
|
|
* - Plans are saved to `.pi/plans/<slug>-<timestamp>.md` (recognizable name + timestamp)
|
|
* - Per-step execution gating: you approve/skip/stop each step before it runs
|
|
* - Progress tracking widget + footer status, persisted across resume
|
|
*/
|
|
|
|
import type { AgentMessage } from "@earendil-works/pi-agent-core";
|
|
import type { TextContent } from "@earendil-works/pi-ai";
|
|
import {
|
|
CONFIG_DIR_NAME,
|
|
type ExtensionAPI,
|
|
type ExtensionContext,
|
|
} from "@earendil-works/pi-coding-agent";
|
|
import { Key } from "@earendil-works/pi-tui";
|
|
import { mkdir, writeFile } from "node:fs/promises";
|
|
import { join } from "node:path";
|
|
import { Type, type Static } from "typebox";
|
|
import { isSafeCommand } from "./utils.ts";
|
|
|
|
// Tools
|
|
const PLAN_MODE_TOOLS = ["read", "bash", "grep", "find", "ls", "questionnaire", "submit_plan"];
|
|
const NORMAL_MODE_TOOLS = ["read", "bash", "edit", "write"];
|
|
const PLAN_MODE_DISABLED_TOOLS = new Set<string>(["edit", "write"]);
|
|
const PLAN_MANAGED_TOOLS = new Set<string>([...PLAN_MODE_TOOLS, ...NORMAL_MODE_TOOLS]);
|
|
|
|
type Phase = "planning" | "awaiting-approval" | "executing";
|
|
|
|
interface PlanStep {
|
|
text: string;
|
|
completed?: boolean;
|
|
skipped?: boolean;
|
|
}
|
|
|
|
interface StoredPlan {
|
|
title: string;
|
|
createdAt: string;
|
|
steps: PlanStep[];
|
|
risks?: string[];
|
|
files?: string[];
|
|
}
|
|
|
|
interface PlanModeState {
|
|
enabled: boolean;
|
|
phase?: Phase;
|
|
plan?: StoredPlan;
|
|
currentStepIndex?: number;
|
|
planFilePath?: string;
|
|
toolsBeforePlanMode?: string[];
|
|
}
|
|
|
|
// ---- submit_plan tool schema ----
|
|
const submitPlanSchema = Type.Object({
|
|
title: Type.String({
|
|
description: "A short, recognizable name for this plan (e.g. 'Add OAuth login').",
|
|
}),
|
|
steps: Type.Array(
|
|
Type.Object({
|
|
description: Type.String({ description: "A single, concrete, actionable step." }),
|
|
}),
|
|
{ description: "Ordered list of steps to implement the plan." },
|
|
),
|
|
risks: Type.Optional(Type.Array(Type.String(), { description: "Potential risks, caveats, or open questions." })),
|
|
files: Type.Optional(Type.Array(Type.String(), { description: "Files likely to be created or modified." })),
|
|
});
|
|
export type SubmitPlanInput = Static<typeof submitPlanSchema>;
|
|
|
|
// ---- file persistence helpers ----
|
|
function slugify(title: string): string {
|
|
const slug = title
|
|
.toLowerCase()
|
|
.replace(/[^a-z0-9]+/g, "-")
|
|
.replace(/^-+|-+$/g, "")
|
|
.slice(0, 50)
|
|
.replace(/-+$/g, "");
|
|
return slug || "plan";
|
|
}
|
|
|
|
function fileTimestamp(): string {
|
|
const d = new Date();
|
|
const p = (n: number) => String(n).padStart(2, "0");
|
|
return `${d.getFullYear()}${p(d.getMonth() + 1)}${p(d.getDate())}-${p(d.getHours())}${p(d.getMinutes())}${p(d.getSeconds())}`;
|
|
}
|
|
|
|
function renderPlanMarkdown(plan: StoredPlan): string {
|
|
const lines: string[] = [];
|
|
lines.push(`# ${plan.title}`);
|
|
lines.push("");
|
|
lines.push(`_Created: ${plan.createdAt}_`);
|
|
lines.push("");
|
|
lines.push("## Steps");
|
|
lines.push("");
|
|
plan.steps.forEach((s, i) => {
|
|
const box = s.completed ? "x" : " ";
|
|
const suffix = s.skipped ? " _(skipped)_" : "";
|
|
lines.push(`${i + 1}. [${box}] ${s.text}${suffix}`);
|
|
});
|
|
if (plan.risks?.length) {
|
|
lines.push("");
|
|
lines.push("## Risks");
|
|
lines.push("");
|
|
for (const r of plan.risks) lines.push(`- ${r}`);
|
|
}
|
|
if (plan.files?.length) {
|
|
lines.push("");
|
|
lines.push("## Files");
|
|
lines.push("");
|
|
for (const f of plan.files) lines.push(`- ${f}`);
|
|
}
|
|
lines.push("");
|
|
return lines.join("\n");
|
|
}
|
|
|
|
export default function planModeExtension(pi: ExtensionAPI): void {
|
|
let planModeEnabled = false;
|
|
let phase: Phase | undefined;
|
|
let plan: StoredPlan | undefined;
|
|
let currentStepIndex = 0;
|
|
let awaitingStep = false;
|
|
let planFilePath: string | undefined;
|
|
let toolsBeforePlanMode: string[] | undefined;
|
|
|
|
pi.registerFlag("plan", {
|
|
description: "Start in plan mode (read-only exploration)",
|
|
type: "boolean",
|
|
default: false,
|
|
});
|
|
|
|
// ---- tool management ----
|
|
function uniqueToolNames(toolNames: string[]): string[] {
|
|
return [...new Set(toolNames)];
|
|
}
|
|
function getPlanModeTools(activeToolNames: string[]): string[] {
|
|
return uniqueToolNames([
|
|
...activeToolNames.filter((name) => !PLAN_MODE_DISABLED_TOOLS.has(name)),
|
|
...PLAN_MODE_TOOLS,
|
|
]);
|
|
}
|
|
function getNormalModeTools(activeToolNames: string[]): string[] {
|
|
return uniqueToolNames([
|
|
...NORMAL_MODE_TOOLS,
|
|
...activeToolNames.filter((name) => !PLAN_MANAGED_TOOLS.has(name)),
|
|
]);
|
|
}
|
|
function enablePlanModeTools(): void {
|
|
if (toolsBeforePlanMode === undefined) {
|
|
toolsBeforePlanMode = pi.getActiveTools();
|
|
}
|
|
pi.setActiveTools(getPlanModeTools(toolsBeforePlanMode));
|
|
}
|
|
function restoreNormalModeTools(): void {
|
|
pi.setActiveTools(toolsBeforePlanMode ?? getNormalModeTools(pi.getActiveTools()));
|
|
toolsBeforePlanMode = undefined;
|
|
}
|
|
|
|
// ---- state ----
|
|
function persistState(): void {
|
|
pi.appendEntry("plan-mode", {
|
|
enabled: planModeEnabled,
|
|
phase,
|
|
plan,
|
|
currentStepIndex,
|
|
planFilePath,
|
|
toolsBeforePlanMode,
|
|
} satisfies PlanModeState);
|
|
}
|
|
|
|
async function persistPlanFile(): Promise<void> {
|
|
if (!planFilePath || !plan) return;
|
|
try {
|
|
await writeFile(planFilePath, renderPlanMarkdown(plan), "utf8");
|
|
} catch {
|
|
// ignore write failures
|
|
}
|
|
}
|
|
|
|
async function savePlanFile(ctx: ExtensionContext, p: StoredPlan): Promise<string> {
|
|
const dir = join(ctx.cwd, CONFIG_DIR_NAME, "plans");
|
|
await mkdir(dir, { recursive: true });
|
|
const path = join(dir, `${slugify(p.title)}-${fileTimestamp()}.md`);
|
|
await writeFile(path, renderPlanMarkdown(p), "utf8");
|
|
return path;
|
|
}
|
|
|
|
// ---- UI ----
|
|
function updateStatus(ctx: ExtensionContext): void {
|
|
if (phase === "executing" && plan) {
|
|
const done = plan.steps.filter((s) => s.completed).length;
|
|
ctx.ui.setStatus("plan-mode", ctx.ui.theme.fg("accent", `📋 ${done}/${plan.steps.length}`));
|
|
} else if (planModeEnabled) {
|
|
ctx.ui.setStatus("plan-mode", ctx.ui.theme.fg("warning", "⏸ plan"));
|
|
} else {
|
|
ctx.ui.setStatus("plan-mode", undefined);
|
|
}
|
|
|
|
if ((phase === "executing" || phase === "awaiting-approval") && plan) {
|
|
const lines = [ctx.ui.theme.fg("accent", plan.title)];
|
|
plan.steps.forEach((s, i) => {
|
|
let mark: string;
|
|
let text = s.text;
|
|
if (s.completed) {
|
|
mark = ctx.ui.theme.fg("success", "☑ ");
|
|
text = ctx.ui.theme.fg("muted", ctx.ui.theme.strikethrough(text));
|
|
} else if (s.skipped) {
|
|
mark = ctx.ui.theme.fg("muted", "⊘ ");
|
|
text = ctx.ui.theme.fg("muted", text);
|
|
} else if (phase === "executing" && i === currentStepIndex) {
|
|
mark = ctx.ui.theme.fg("accent", "▶ ");
|
|
} else {
|
|
mark = ctx.ui.theme.fg("muted", "☐ ");
|
|
}
|
|
lines.push(`${mark}${text}`);
|
|
});
|
|
ctx.ui.setWidget("plan-todos", lines);
|
|
} else {
|
|
ctx.ui.setWidget("plan-todos", undefined);
|
|
}
|
|
}
|
|
|
|
// ---- mode toggle ----
|
|
function togglePlanMode(ctx: ExtensionContext): void {
|
|
planModeEnabled = !planModeEnabled;
|
|
phase = planModeEnabled ? "planning" : undefined;
|
|
plan = undefined;
|
|
planFilePath = undefined;
|
|
currentStepIndex = 0;
|
|
awaitingStep = false;
|
|
|
|
if (planModeEnabled) {
|
|
enablePlanModeTools();
|
|
ctx.ui.notify("Plan mode enabled. Built-in write tools disabled. Ask for a plan, then submit_plan.");
|
|
} else {
|
|
restoreNormalModeTools();
|
|
ctx.ui.notify("Plan mode disabled. Full access restored.");
|
|
}
|
|
updateStatus(ctx);
|
|
persistState();
|
|
}
|
|
|
|
// ---- execution gating ----
|
|
function finishExecution(ctx: ExtensionContext): void {
|
|
if (plan) {
|
|
const summary = plan.steps
|
|
.map((s) => (s.completed ? `~~${s.text}~~` : s.skipped ? `${s.text} (skipped)` : s.text))
|
|
.join("\n");
|
|
pi.sendMessage(
|
|
{ customType: "plan-complete", content: `**Plan "${plan.title}" complete!** ✓\n\n${summary}`, display: true },
|
|
{ triggerTurn: false },
|
|
);
|
|
}
|
|
phase = undefined;
|
|
awaitingStep = false;
|
|
updateStatus(ctx);
|
|
persistState();
|
|
void persistPlanFile();
|
|
}
|
|
|
|
async function gateNext(ctx: ExtensionContext): Promise<void> {
|
|
if (!plan) {
|
|
phase = undefined;
|
|
return;
|
|
}
|
|
if (currentStepIndex >= plan.steps.length) {
|
|
finishExecution(ctx);
|
|
return;
|
|
}
|
|
|
|
const step = plan.steps[currentStepIndex];
|
|
if (!step) {
|
|
finishExecution(ctx);
|
|
return;
|
|
}
|
|
|
|
const choice = await ctx.ui.select(
|
|
`Step ${currentStepIndex + 1}/${plan.steps.length}: ${step.text}`,
|
|
["Execute this step", "Skip this step", "Stop execution"],
|
|
);
|
|
|
|
if (choice?.startsWith("Execute")) {
|
|
awaitingStep = true;
|
|
updateStatus(ctx);
|
|
const stepMsg = `Execute ONLY step ${currentStepIndex + 1} of ${plan.steps.length} from plan "${plan.title}":
|
|
|
|
${step.text}
|
|
|
|
Do just this one step. Do NOT start any other step. Stop when this step is done.`;
|
|
pi.sendMessage(
|
|
{ customType: "plan-step", content: stepMsg, display: true },
|
|
{ triggerTurn: true, deliverAs: "followUp" },
|
|
);
|
|
} else if (choice?.startsWith("Skip")) {
|
|
step.skipped = true;
|
|
currentStepIndex++;
|
|
updateStatus(ctx);
|
|
persistState();
|
|
await persistPlanFile();
|
|
await gateNext(ctx);
|
|
} else {
|
|
ctx.ui.notify("Plan execution stopped.", "info");
|
|
phase = undefined;
|
|
awaitingStep = false;
|
|
updateStatus(ctx);
|
|
persistState();
|
|
}
|
|
}
|
|
|
|
async function handleApproval(ctx: ExtensionContext): Promise<void> {
|
|
if (!plan) {
|
|
phase = undefined;
|
|
return;
|
|
}
|
|
const where = planFilePath ? `\nSaved to: ${planFilePath}` : "";
|
|
const choice = await ctx.ui.select(
|
|
`Plan "${plan.title}" — ${plan.steps.length} steps.${where}\nWhat next?`,
|
|
["Execute step-by-step", "Stay in plan mode", "Refine the plan"],
|
|
);
|
|
|
|
if (choice?.startsWith("Execute")) {
|
|
planModeEnabled = false;
|
|
phase = "executing";
|
|
currentStepIndex = 0;
|
|
awaitingStep = false;
|
|
restoreNormalModeTools();
|
|
updateStatus(ctx);
|
|
persistState();
|
|
await gateNext(ctx);
|
|
} else if (choice === "Refine the plan") {
|
|
phase = "planning";
|
|
const refinement = await ctx.ui.editor("Refine the plan:", "");
|
|
updateStatus(ctx);
|
|
persistState();
|
|
if (refinement?.trim()) {
|
|
pi.sendUserMessage(refinement.trim(), { deliverAs: "followUp" });
|
|
}
|
|
} else {
|
|
phase = "planning";
|
|
updateStatus(ctx);
|
|
persistState();
|
|
}
|
|
}
|
|
|
|
// ---- commands & shortcuts ----
|
|
pi.registerCommand("plan", {
|
|
description: "Toggle plan mode (read-only exploration + structured planning)",
|
|
handler: async (_args, ctx) => togglePlanMode(ctx),
|
|
});
|
|
|
|
pi.registerCommand("plan-status", {
|
|
description: "Show the current plan and progress",
|
|
handler: async (_args, ctx) => {
|
|
if (!plan) {
|
|
ctx.ui.notify("No plan yet. Enable /plan and ask the agent to submit_plan.", "info");
|
|
return;
|
|
}
|
|
const list = plan.steps
|
|
.map((s, i) => `${i + 1}. ${s.completed ? "✓" : s.skipped ? "⊘" : "○"} ${s.text}`)
|
|
.join("\n");
|
|
const loc = planFilePath ? `\n\nFile: ${planFilePath}` : "";
|
|
ctx.ui.notify(`Plan "${plan.title}" (${phase ?? "idle"}):\n${list}${loc}`, "info");
|
|
},
|
|
});
|
|
|
|
pi.registerCommand("plan-continue", {
|
|
description: "Resume gated step-by-step execution of the current plan",
|
|
handler: async (_args, ctx) => {
|
|
if (phase !== "executing" || !plan) {
|
|
ctx.ui.notify("No plan execution in progress.", "info");
|
|
return;
|
|
}
|
|
if (awaitingStep) {
|
|
ctx.ui.notify("A step is already running.", "info");
|
|
return;
|
|
}
|
|
await gateNext(ctx);
|
|
},
|
|
});
|
|
|
|
pi.registerShortcut(Key.ctrlAlt("p"), {
|
|
description: "Toggle plan mode",
|
|
handler: async (ctx) => togglePlanMode(ctx),
|
|
});
|
|
|
|
// ---- submit_plan tool ----
|
|
pi.registerTool({
|
|
name: "submit_plan",
|
|
label: "Submit Plan",
|
|
description:
|
|
"Submit a structured implementation plan for user approval. ONLY available in plan mode. " +
|
|
"Call this once analysis is finished and you are ready to propose the work. After calling it, stop and wait for the user.",
|
|
promptSnippet: "Submit a structured plan (title + numbered steps) for approval while in plan mode",
|
|
promptGuidelines: [
|
|
"Use submit_plan to deliver your plan when in plan mode instead of writing a free-form 'Plan:' section.",
|
|
"After calling submit_plan, do not take further action; wait for the user to approve, refine, or execute.",
|
|
],
|
|
parameters: submitPlanSchema,
|
|
async execute(_toolCallId, params: SubmitPlanInput, _signal, _onUpdate, ctx) {
|
|
if (!planModeEnabled) {
|
|
return {
|
|
content: [
|
|
{ type: "text", text: "submit_plan is only available in plan mode. Enable it with /plan first." },
|
|
],
|
|
isError: true,
|
|
details: {},
|
|
};
|
|
}
|
|
|
|
plan = {
|
|
title: params.title.trim() || "Untitled plan",
|
|
createdAt: new Date().toISOString(),
|
|
steps: params.steps
|
|
.map((s) => ({ text: s.description.trim(), completed: false }))
|
|
.filter((s) => s.text.length > 0),
|
|
risks: params.risks?.map((r) => r.trim()).filter(Boolean),
|
|
files: params.files?.map((f) => f.trim()).filter(Boolean),
|
|
};
|
|
currentStepIndex = 0;
|
|
awaitingStep = false;
|
|
phase = "awaiting-approval";
|
|
|
|
try {
|
|
planFilePath = await savePlanFile(ctx, plan);
|
|
} catch {
|
|
planFilePath = undefined;
|
|
}
|
|
persistState();
|
|
updateStatus(ctx);
|
|
|
|
const loc = planFilePath ? ` and saved to ${planFilePath}` : "";
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: `Plan "${plan.title}" with ${plan.steps.length} step(s) submitted${loc}. Stop here and wait for the user to approve, refine, or execute.`,
|
|
},
|
|
],
|
|
details: { title: plan.title, steps: plan.steps.length, path: planFilePath ?? null },
|
|
terminate: true,
|
|
};
|
|
},
|
|
});
|
|
|
|
// ---- block destructive bash in plan mode ----
|
|
pi.on("tool_call", async (event) => {
|
|
if (!planModeEnabled || event.toolName !== "bash") return;
|
|
const command = event.input.command as string;
|
|
if (!isSafeCommand(command)) {
|
|
return {
|
|
block: true,
|
|
reason: `Plan mode: command blocked (not allowlisted). Use /plan to disable plan mode first.\nCommand: ${command}`,
|
|
};
|
|
}
|
|
});
|
|
|
|
// ---- strip stale plan-mode context when not in plan mode ----
|
|
pi.on("context", async (event) => {
|
|
if (planModeEnabled) return;
|
|
return {
|
|
messages: event.messages.filter((m) => {
|
|
const msg = m as AgentMessage & { customType?: string };
|
|
if (msg.customType === "plan-mode-context") return false;
|
|
if (msg.role !== "user") return true;
|
|
const content = msg.content;
|
|
if (typeof content === "string") return !content.includes("[PLAN MODE ACTIVE]");
|
|
if (Array.isArray(content)) {
|
|
return !content.some(
|
|
(c) => c.type === "text" && (c as TextContent).text?.includes("[PLAN MODE ACTIVE]"),
|
|
);
|
|
}
|
|
return true;
|
|
}),
|
|
};
|
|
});
|
|
|
|
// ---- inject plan-mode instructions ----
|
|
pi.on("before_agent_start", async () => {
|
|
if (!planModeEnabled) return;
|
|
return {
|
|
message: {
|
|
customType: "plan-mode-context",
|
|
content: `[PLAN MODE ACTIVE]
|
|
You are in plan mode - a read-only exploration mode for safe code analysis.
|
|
|
|
Restrictions:
|
|
- Built-in edit and write tools are disabled
|
|
- Bash is restricted to an allowlist of read-only commands
|
|
- Do NOT attempt to make changes - only investigate and plan
|
|
|
|
Explore the code, ask clarifying questions if needed, then call the "submit_plan" tool with:
|
|
- title: a short, recognizable name for the plan
|
|
- steps: an ordered list of concrete, actionable steps
|
|
- risks: optional caveats or open questions
|
|
- files: optional files you expect to create or modify
|
|
|
|
Do NOT write a free-form "Plan:" section - use the submit_plan tool instead.
|
|
After submitting, stop and wait for the user.`,
|
|
display: false,
|
|
},
|
|
};
|
|
});
|
|
|
|
// ---- drive approval + gated execution when the agent settles ----
|
|
pi.on("agent_settled", async (_event, ctx) => {
|
|
if (!ctx.hasUI) return;
|
|
|
|
if (phase === "awaiting-approval") {
|
|
await handleApproval(ctx);
|
|
return;
|
|
}
|
|
|
|
if (phase === "executing" && awaitingStep) {
|
|
awaitingStep = false;
|
|
const step = plan?.steps[currentStepIndex];
|
|
if (step) step.completed = true;
|
|
currentStepIndex++;
|
|
updateStatus(ctx);
|
|
persistState();
|
|
await persistPlanFile();
|
|
await gateNext(ctx);
|
|
}
|
|
});
|
|
|
|
// ---- restore state on session start / resume ----
|
|
pi.on("session_start", async (_event, ctx) => {
|
|
if (pi.getFlag("plan") === true) {
|
|
planModeEnabled = true;
|
|
phase = "planning";
|
|
}
|
|
|
|
const entries = ctx.sessionManager.getEntries();
|
|
const planModeEntry = entries
|
|
.filter((e: { type: string; customType?: string }) => e.type === "custom" && e.customType === "plan-mode")
|
|
.pop() as { data?: PlanModeState } | undefined;
|
|
|
|
if (planModeEntry?.data) {
|
|
const d = planModeEntry.data;
|
|
planModeEnabled = d.enabled ?? planModeEnabled;
|
|
phase = d.phase ?? phase;
|
|
plan = d.plan ?? plan;
|
|
currentStepIndex = d.currentStepIndex ?? currentStepIndex;
|
|
planFilePath = d.planFilePath ?? planFilePath;
|
|
toolsBeforePlanMode = d.toolsBeforePlanMode ?? toolsBeforePlanMode;
|
|
awaitingStep = false; // never resume mid-step automatically
|
|
}
|
|
|
|
if (planModeEnabled) {
|
|
enablePlanModeTools();
|
|
}
|
|
updateStatus(ctx);
|
|
|
|
if (phase === "executing" && plan && ctx.hasUI) {
|
|
ctx.ui.notify(`Plan "${plan.title}" execution paused. Use /plan-continue to resume.`, "info");
|
|
}
|
|
});
|
|
}
|