410 lines
12 KiB
TypeScript
410 lines
12 KiB
TypeScript
import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
|
|
import { extractTodoItems, isSafeCommand, markCompletedSteps, type TodoItem } from "./utils.js";
|
|
|
|
const PLAN_FLAG = "plan-mode";
|
|
const PLAN_TODOS_WIDGET = "plan-todos";
|
|
const PLAN_CONTEXT_TYPE = "plan-mode-context";
|
|
const PLAN_EXECUTION_CONTEXT_TYPE = "plan-execution-context";
|
|
const PLAN_STATE_ENTRY = "plan-mode-state";
|
|
const PREFERRED_PLAN_TOOLS = ["read", "bash", "grep", "find", "ls", "lsp_diagnostics", "subagent"];
|
|
|
|
type PlanModeState = {
|
|
enabled: boolean;
|
|
executing: boolean;
|
|
savedTools: string[] | null;
|
|
todos: TodoItem[];
|
|
};
|
|
|
|
function getAssistantText(message: any): string {
|
|
if (!message || message.role !== "assistant" || !Array.isArray(message.content)) return "";
|
|
return message.content
|
|
.filter((part: any) => part?.type === "text" && typeof part.text === "string")
|
|
.map((part: any) => part.text)
|
|
.join("\n")
|
|
.trim();
|
|
}
|
|
|
|
export default function planModeExtension(pi: ExtensionAPI): void {
|
|
let planModeEnabled = false;
|
|
let executionMode = false;
|
|
let savedTools: string[] | null = null;
|
|
let todoItems: TodoItem[] = [];
|
|
|
|
function persistState() {
|
|
const state: PlanModeState = {
|
|
enabled: planModeEnabled,
|
|
executing: executionMode,
|
|
savedTools,
|
|
todos: todoItems,
|
|
};
|
|
pi.appendEntry(PLAN_STATE_ENTRY, state);
|
|
}
|
|
|
|
function getPlanTools(): string[] {
|
|
const allToolNames = new Set(pi.getAllTools().map((tool) => tool.name));
|
|
const tools = new Set<string>();
|
|
|
|
for (const name of PREFERRED_PLAN_TOOLS) {
|
|
if (allToolNames.has(name)) tools.add(name);
|
|
}
|
|
|
|
if (savedTools) {
|
|
for (const name of savedTools) {
|
|
if (allToolNames.has(name) && PREFERRED_PLAN_TOOLS.includes(name)) tools.add(name);
|
|
}
|
|
}
|
|
|
|
return Array.from(tools);
|
|
}
|
|
|
|
function restoreSavedTools() {
|
|
if (savedTools && savedTools.length > 0) {
|
|
pi.setActiveTools(savedTools);
|
|
}
|
|
}
|
|
|
|
function updateUi(ctx: ExtensionContext) {
|
|
if (executionMode && todoItems.length > 0) {
|
|
const completed = todoItems.filter((item) => item.completed).length;
|
|
ctx.ui.setStatus(PLAN_FLAG, `📋 ${completed}/${todoItems.length}`);
|
|
} else if (planModeEnabled) {
|
|
ctx.ui.setStatus(PLAN_FLAG, "⏸ PLAN");
|
|
} else {
|
|
ctx.ui.setStatus(PLAN_FLAG, undefined);
|
|
}
|
|
|
|
if ((planModeEnabled || executionMode) && todoItems.length > 0) {
|
|
ctx.ui.setWidget(
|
|
PLAN_TODOS_WIDGET,
|
|
todoItems.map((item) => `${item.completed ? "☑" : "☐"} ${item.step}. ${item.text}`),
|
|
);
|
|
} else {
|
|
ctx.ui.setWidget(PLAN_TODOS_WIDGET, undefined);
|
|
}
|
|
}
|
|
|
|
function enablePlanMode(ctx: ExtensionContext) {
|
|
if (!savedTools) savedTools = pi.getActiveTools();
|
|
planModeEnabled = true;
|
|
executionMode = false;
|
|
pi.setActiveTools(getPlanTools());
|
|
updateUi(ctx);
|
|
persistState();
|
|
ctx.ui.notify(`Plan mode enabled. Active tools: ${getPlanTools().join(", ")}`, "info");
|
|
}
|
|
|
|
function disablePlanMode(ctx: ExtensionContext, notify = true) {
|
|
planModeEnabled = false;
|
|
executionMode = false;
|
|
restoreSavedTools();
|
|
updateUi(ctx);
|
|
persistState();
|
|
if (notify) ctx.ui.notify("Plan mode disabled. Previous tools restored.", "info");
|
|
}
|
|
|
|
function clearPlan(ctx: ExtensionContext) {
|
|
executionMode = false;
|
|
todoItems = [];
|
|
updateUi(ctx);
|
|
persistState();
|
|
ctx.ui.notify("Stored plan cleared.", "info");
|
|
}
|
|
|
|
function showStatus(ctx: ExtensionContext) {
|
|
if (!planModeEnabled && !executionMode && todoItems.length === 0) {
|
|
ctx.ui.notify("Plan mode is off and no stored plan exists.", "info");
|
|
return;
|
|
}
|
|
|
|
const lines = [
|
|
`plan mode: ${planModeEnabled ? "on" : "off"}`,
|
|
`execution: ${executionMode ? "on" : "off"}`,
|
|
`stored steps: ${todoItems.length}`,
|
|
];
|
|
if (todoItems.length > 0) {
|
|
for (const item of todoItems) lines.push(`${item.step}. ${item.completed ? "✓" : "○"} ${item.text}`);
|
|
}
|
|
ctx.ui.notify(lines.join("\n"), "info");
|
|
}
|
|
|
|
function getLatestPlannedAssistantText(ctx: ExtensionContext): string {
|
|
const branch = ctx.sessionManager.getBranch();
|
|
for (let i = branch.length - 1; i >= 0; i--) {
|
|
const entry: any = branch[i];
|
|
if (entry?.type !== "message" || !entry.message) continue;
|
|
const text = getAssistantText(entry.message);
|
|
if (text && extractTodoItems(text).length > 0) return text;
|
|
}
|
|
return "";
|
|
}
|
|
|
|
function buildExecutionPrompt(ctx: ExtensionContext, fresh: boolean): string {
|
|
const remaining = todoItems.filter((item) => !item.completed);
|
|
const planLines = todoItems.map((item) => `${item.step}. ${item.text}${item.completed ? " [already done]" : ""}`);
|
|
const latestPlanText = getLatestPlannedAssistantText(ctx);
|
|
const sections = [
|
|
fresh
|
|
? "You are executing an approved plan in a fresh session. The earlier planning conversation is not available, so rely only on the context below."
|
|
: "Execute the approved plan in this session.",
|
|
"## Approved Plan",
|
|
planLines.join("\n"),
|
|
];
|
|
|
|
if (remaining.length > 0) {
|
|
sections.push("## Remaining Steps", remaining.map((item) => `${item.step}. ${item.text}`).join("\n"));
|
|
}
|
|
|
|
if (latestPlanText) {
|
|
sections.push("## Planner Notes", latestPlanText);
|
|
}
|
|
|
|
sections.push(
|
|
"## Execution Instructions",
|
|
[
|
|
"Inspect files before editing.",
|
|
"If the approved plan needs adjustment, explain why before making a larger change.",
|
|
"Complete the remaining steps in order.",
|
|
"After completing a step, include a [DONE:n] marker in your response.",
|
|
].join("\n"),
|
|
);
|
|
|
|
return sections.join("\n\n");
|
|
}
|
|
|
|
function startExecution(ctx: ExtensionContext) {
|
|
if (todoItems.length === 0) {
|
|
ctx.ui.notify("No stored plan found. Create one in /plan mode first.", "warning");
|
|
return;
|
|
}
|
|
|
|
planModeEnabled = false;
|
|
executionMode = true;
|
|
restoreSavedTools();
|
|
updateUi(ctx);
|
|
persistState();
|
|
ctx.ui.notify("Executing stored plan with full tools restored.", "info");
|
|
pi.sendUserMessage(buildExecutionPrompt(ctx, false));
|
|
}
|
|
|
|
async function startFreshExecution(ctx: any) {
|
|
if (todoItems.length === 0) {
|
|
ctx.ui.notify("No stored plan found. Create one in /plan mode first.", "warning");
|
|
return;
|
|
}
|
|
|
|
await ctx.waitForIdle();
|
|
|
|
const executionPrompt = buildExecutionPrompt(ctx, true);
|
|
const executionTools = savedTools && savedTools.length > 0 ? [...savedTools] : [...pi.getActiveTools()];
|
|
const carriedTodos = todoItems.map((item) => ({ ...item }));
|
|
const parentSession = ctx.sessionManager.getSessionFile();
|
|
|
|
const result = await ctx.newSession({
|
|
parentSession,
|
|
setup: async (sm: any) => {
|
|
sm.appendCustomEntry(PLAN_STATE_ENTRY, {
|
|
enabled: false,
|
|
executing: true,
|
|
savedTools: executionTools,
|
|
todos: carriedTodos,
|
|
});
|
|
},
|
|
});
|
|
|
|
if (result.cancelled) {
|
|
ctx.ui.notify("Fresh execution session cancelled.", "info");
|
|
return;
|
|
}
|
|
|
|
pi.sendUserMessage(executionPrompt);
|
|
}
|
|
|
|
pi.registerFlag("plan", {
|
|
description: "Start in read-only plan mode",
|
|
type: "boolean",
|
|
default: false,
|
|
});
|
|
|
|
pi.registerCommand("plan", {
|
|
description: "Toggle plan mode, or use /plan on|off|execute [fresh]|status|clear",
|
|
handler: async (args, ctx) => {
|
|
const action = args.trim().toLowerCase();
|
|
switch (action) {
|
|
case "":
|
|
case "toggle":
|
|
if (planModeEnabled) disablePlanMode(ctx);
|
|
else enablePlanMode(ctx);
|
|
return;
|
|
case "on":
|
|
enablePlanMode(ctx);
|
|
return;
|
|
case "off":
|
|
disablePlanMode(ctx);
|
|
return;
|
|
case "execute":
|
|
startExecution(ctx);
|
|
return;
|
|
case "execute fresh":
|
|
await startFreshExecution(ctx);
|
|
return;
|
|
case "status":
|
|
showStatus(ctx);
|
|
return;
|
|
case "clear":
|
|
clearPlan(ctx);
|
|
return;
|
|
default:
|
|
ctx.ui.notify("Usage: /plan [on|off|execute|execute fresh|status|clear]", "warning");
|
|
}
|
|
},
|
|
});
|
|
|
|
pi.registerShortcut("ctrl+alt+p", {
|
|
description: "Toggle plan mode on/off",
|
|
handler: async (ctx) => {
|
|
if (planModeEnabled) disablePlanMode(ctx);
|
|
else enablePlanMode(ctx);
|
|
},
|
|
});
|
|
|
|
pi.on("tool_call", async (event) => {
|
|
if (!planModeEnabled) return;
|
|
|
|
if (event.toolName === "edit" || event.toolName === "write") {
|
|
return { block: true, reason: "Plan mode is read-only. Disable it before modifying files." };
|
|
}
|
|
|
|
if (event.toolName === "bash") {
|
|
const command = typeof event.input?.command === "string" ? event.input.command : "";
|
|
if (!isSafeCommand(command)) {
|
|
return {
|
|
block: true,
|
|
reason: `Plan mode blocked a non read-only bash command: ${command}`,
|
|
};
|
|
}
|
|
}
|
|
|
|
});
|
|
|
|
pi.on("context", async (event) => {
|
|
if (planModeEnabled || executionMode) return;
|
|
return {
|
|
messages: event.messages.filter((message: any) => {
|
|
if (message?.customType === PLAN_CONTEXT_TYPE) return false;
|
|
if (message?.customType === PLAN_EXECUTION_CONTEXT_TYPE) return false;
|
|
return true;
|
|
}),
|
|
};
|
|
});
|
|
|
|
pi.on("before_agent_start", async () => {
|
|
if (planModeEnabled) {
|
|
return {
|
|
message: {
|
|
customType: PLAN_CONTEXT_TYPE,
|
|
display: false,
|
|
content: `[PLAN MODE ACTIVE]\nYou are in read-only planning mode.\n\nRules:\n- Do not modify files.\n- Do not use edit or write.\n- Use only read-only investigation.\n- You may use subagent when it helps planning. Prefer read-only helpers like librarian, reviewer, or uiux-designer when appropriate.\n- If requirements are unclear, ask clarifying questions before planning.\n- Produce a numbered plan under a \"Plan:\" header.\n- Call out risks, dependencies, and validation steps.`,
|
|
},
|
|
};
|
|
}
|
|
|
|
if (executionMode && todoItems.length > 0) {
|
|
const remaining = todoItems.filter((item) => !item.completed);
|
|
return {
|
|
message: {
|
|
customType: PLAN_EXECUTION_CONTEXT_TYPE,
|
|
display: false,
|
|
content:
|
|
`[EXECUTING APPROVED PLAN]\nComplete the remaining plan steps in order.\n\nRemaining steps:\n${remaining.map((item) => `${item.step}. ${item.text}`).join("\n")}\n\nAfter completing a step, include [DONE:n] in your response.`,
|
|
},
|
|
};
|
|
}
|
|
});
|
|
|
|
pi.on("turn_end", async (event, ctx) => {
|
|
if (!executionMode || todoItems.length === 0) return;
|
|
const text = getAssistantText(event.message);
|
|
if (!text) return;
|
|
if (markCompletedSteps(text, todoItems) > 0) {
|
|
updateUi(ctx);
|
|
persistState();
|
|
}
|
|
if (todoItems.length > 0 && todoItems.every((item) => item.completed)) {
|
|
executionMode = false;
|
|
updateUi(ctx);
|
|
persistState();
|
|
ctx.ui.notify("Plan execution completed.", "success");
|
|
}
|
|
});
|
|
|
|
pi.on("agent_end", async (event, ctx) => {
|
|
if (!planModeEnabled) return;
|
|
const lastAssistant = [...event.messages].reverse().find((message: any) => message?.role === "assistant");
|
|
const text = getAssistantText(lastAssistant);
|
|
if (!text) return;
|
|
|
|
const extracted = extractTodoItems(text);
|
|
if (extracted.length === 0) return;
|
|
|
|
todoItems = extracted;
|
|
updateUi(ctx);
|
|
persistState();
|
|
|
|
if (!ctx.hasUI) return;
|
|
|
|
const choice = await ctx.ui.select("Plan ready - what next?", [
|
|
"Execute the plan now",
|
|
"Execute in a fresh session",
|
|
"Send custom feedback / continue discussion",
|
|
"Stay in plan mode",
|
|
]);
|
|
|
|
if (choice === "Execute the plan now") {
|
|
startExecution(ctx);
|
|
return;
|
|
}
|
|
|
|
if (choice === "Execute in a fresh session") {
|
|
await startFreshExecution(ctx);
|
|
return;
|
|
}
|
|
|
|
if (choice === "Send custom feedback / continue discussion") {
|
|
const feedback = await ctx.ui.editor(
|
|
"What should Pi do next?",
|
|
"Refine the plan by ...",
|
|
);
|
|
if (feedback?.trim()) {
|
|
pi.sendUserMessage(feedback.trim());
|
|
}
|
|
}
|
|
});
|
|
|
|
pi.on("session_start", async (_event, ctx) => {
|
|
const branch = ctx.sessionManager.getBranch();
|
|
const stateEntry = branch
|
|
.filter((entry: any) => entry.type === "custom" && entry.customType === PLAN_STATE_ENTRY)
|
|
.pop() as { data?: PlanModeState } | undefined;
|
|
|
|
if (stateEntry?.data) {
|
|
planModeEnabled = stateEntry.data.enabled ?? false;
|
|
executionMode = stateEntry.data.executing ?? false;
|
|
savedTools = stateEntry.data.savedTools ?? null;
|
|
todoItems = stateEntry.data.todos ?? [];
|
|
} else if (pi.getFlag("--plan") === true || pi.getFlag("plan") === true) {
|
|
planModeEnabled = true;
|
|
executionMode = false;
|
|
savedTools = pi.getActiveTools();
|
|
}
|
|
|
|
if (!savedTools) savedTools = pi.getActiveTools();
|
|
|
|
if (planModeEnabled) {
|
|
pi.setActiveTools(getPlanTools());
|
|
} else if (executionMode) {
|
|
restoreSavedTools();
|
|
}
|
|
|
|
updateUi(ctx);
|
|
});
|
|
}
|