@dreki-gg/pi-plan-mode 0.2.0 → 0.3.1

package/extensions/plan-mode/index.ts

@@ -1,664 +1,579 @@
-import type { ExtensionAPI, ExtensionContext } from '@mariozechner/pi-coding-agent';
-import { Key, type AutocompleteItem } from '@mariozechner/pi-tui';
+/**
+ * Plan Mode Extension
+ *
+ * Two-phase workflow:
+ *   1. PLAN phase  — read-only tools (+ edit/write for .plans/ only) + medium thinking
+ *                    Planner analyzes codebase, asks questions, writes PLAN.md + START-PROMPT.md
+ *   2. EXECUTE phase — full tools + low thinking, clean context from START-PROMPT.md
+ *                      Executor works through the plan step by step with [DONE:n] tracking
+ *
+ * Plans live in `.plans/<kebab-name>/PLAN.md` with a `START-PROMPT.md` sibling for clean handoff.
+ *
+ * Commands:
+ *   /plan [prompt]  — enter plan mode (optionally with a starting prompt)
+ *   /todos          — show current plan progress
+ *   Ctrl+Alt+P      — toggle plan mode (shortcut)
+ *
+ * Flag:
+ *   --plan          — start session in plan mode
+ */
+
+import type { AgentMessage } from '@earendil-works/pi-agent-core';
+import type { AssistantMessage, TextContent } from '@earendil-works/pi-ai';
+import type { ExtensionAPI, ExtensionContext } from '@earendil-works/pi-coding-agent';
+import { Key } from '@earendil-works/pi-tui';
 import {
   extractTodoItems,
-  formatTodoList,
   isSafeCommand,
   markCompletedSteps,
   type TodoItem,
 } from './utils.js';
 
-type WorkflowPhase = 'off' | 'planning' | 'plan-files' | 'executing';
-
-interface PersistedState {
-  phase: WorkflowPhase;
-  todos: TodoItem[];
-}
-
-interface BeforeAgentStartCompatEvent {
-  systemPromptOptions?: {
-    selectedTools?: string[];
-  };
-}
-
-const STATE_ENTRY = 'plan-mode-state';
-const CLEAR_VALUES = new Set(['', 'off', 'none', 'disable', 'exit']);
-const PLANNING_TOOL_PATTERNS = [
+// ── Tool sets ────────────────────────────────────────────────────────────────
+// Plan phase: read-only + edit/write (for .plans/ files only, enforced by prompt)
+const PLAN_TOOLS = [
   'read',
   'bash',
   'grep',
   'find',
   'ls',
+  'edit',
+  'write',
   'questionnaire',
-  'lsp',
-] as const;
-const PLAN_FILE_TOOL_PATTERNS = [...PLANNING_TOOL_PATTERNS, 'edit', 'write'] as const;
-
-function isAssistantMessage(message: unknown): message is {
-  role: 'assistant';
-  content: Array<{ type: string; text?: string }>;
-} {
-  return (
-    typeof message === 'object' &&
-    message !== null &&
-    'role' in message &&
-    (message as { role?: unknown }).role === 'assistant' &&
-    'content' in message &&
-    Array.isArray((message as { content?: unknown }).content)
-  );
-}
+  'search_skills',
+];
+const EXEC_TOOLS = ['read', 'bash', 'edit', 'write', 'search_skills'];
 
-function getMessageText(message: unknown): string {
-  if (!isAssistantMessage(message)) return '';
-  return message.content
-    .filter(
-      (block): block is { type: 'text'; text: string } =>
-        block.type === 'text' && typeof block.text === 'string',
-    )
-    .map((block) => block.text)
-    .join('\n');
-}
+// ── Model + thinking presets ─────────────────────────────────────────────────
+const PLAN_MODEL = { provider: 'anthropic', id: 'claude-opus-4-6' } as const;
+const PLAN_THINKING = 'medium' as const;
 
-function matchesToolPattern(name: string, pattern: string): boolean {
-  if (pattern.endsWith('*')) return name.startsWith(pattern.slice(0, -1));
-  return name === pattern;
-}
+const EXEC_MODEL = { provider: 'openai', id: 'gpt-5.5' } as const;
+const EXEC_THINKING = 'low' as const;
 
-function resolveToolNames(allToolNames: string[], patterns: readonly string[]): string[] {
-  return allToolNames.filter((name) =>
-    patterns.some(
-      (pattern) =>
-        matchesToolPattern(name, pattern) ||
-        (pattern === 'context7_*' && name.startsWith('context7_')),
-    ),
-  );
-}
+type ThinkingLevel = 'off' | 'minimal' | 'low' | 'medium' | 'high' | 'xhigh';
 
-function findSavedState(ctx: ExtensionContext): PersistedState | undefined {
-  const branchEntries = ctx.sessionManager.getBranch?.() ?? ctx.sessionManager.getEntries?.() ?? [];
-  let lastState: PersistedState | undefined;
-
-  for (const entry of branchEntries) {
-    if (entry.type !== 'custom' || entry.customType !== STATE_ENTRY) continue;
-    const data = entry.data as PersistedState | undefined;
-    if (data && typeof data.phase === 'string' && Array.isArray(data.todos)) lastState = data;
-  }
-
-  return lastState;
+// ── Persisted state ──────────────────────────────────────────────────────────
+interface PersistedState {
+  planEnabled: boolean;
+  executing: boolean;
+  planDir: string | undefined;
+  todos: TodoItem[];
 }
 
-function buildPlanningInstructions(
-  ctx: ExtensionContext,
-  activeTools: string[],
-  todos: TodoItem[],
-): string {
-  const sections = [
-    'PLAN MODE ACTIVE.',
-    'Run a Cursor-style planning pass before implementation.',
-    `Enabled tools: ${activeTools.length > 0 ? activeTools.join(', ') : '(none)'}`,
-    'Rules:',
-    '- Stay read-only. Do not edit files or claim to have changed code.',
-    '- Inspect the real codebase before proposing work. Do not guess from filenames.',
-    '- If requirements are underspecified and the questionnaire tool is available, use it to ask 1-5 structured clarifying questions before finalizing the plan.',
-    '- Produce a concrete numbered plan under a `Plan:` header.',
-    '- Include open questions, assumptions, risks, and sequencing when relevant.',
-    '- If the user wants terminology hardened or the design pressure-tested, point them to the domain-model workflow.',
-    '- If the user wants self-contained handoff files, point them to the implementation-plan workflow.',
-  ];
-
-  if (ctx.hasUI) {
-    sections.push(
-      'After the planning response, the extension may offer next-step choices like domain-model review, implementation-plan generation, or execution.',
-    );
-  }
-
-  if (todos.length > 0) {
-    sections.push(`Current plan draft:\n${formatTodoList(todos)}`);
-  }
-
-  return sections.join('\n');
+// ── Helpers ──────────────────────────────────────────────────────────────────
+function isAssistantMessage(m: AgentMessage): m is AssistantMessage {
+  return m.role === 'assistant' && Array.isArray(m.content);
 }
 
-function buildPlanFileInstructions(activeTools: string[], todos: TodoItem[]): string {
-  const sections = [
-    'PLAN FILE AUTHORING MODE ACTIVE.',
-    `Enabled tools: ${activeTools.length > 0 ? activeTools.join(', ') : '(none)'}`,
-    'You may use edit/write in this phase, but only to author planning artifacts such as `*.plan.md`, `CONTEXT.md`, or ADR docs requested by the workflow.',
-    'Do not implement product code in this phase.',
-    'Ground every plan file in the current codebase state.',
-  ];
-
-  if (todos.length > 0) {
-    sections.push(`Current approved plan:\n${formatTodoList(todos)}`);
-  }
-
-  return sections.join('\n');
-}
-
-function buildExecutionInstructions(todos: TodoItem[]): string {
-  const sections = [
-    'PLAN EXECUTION MODE ACTIVE.',
-    'Execute the approved plan in small verified steps.',
-    'After completing a plan step, include a `[DONE:n]` tag in the response so progress can be tracked.',
-  ];
-
-  if (todos.length > 0) {
-    sections.push(`Current checklist:\n${formatTodoList(todos)}`);
-  }
-
-  return sections.join('\n');
+function getTextContent(message: AssistantMessage): string {
+  return message.content
+    .filter((b): b is TextContent => b.type === 'text')
+    .map((b) => b.text)
+    .join('\n');
 }
 
-export default function planModeExtension(pi: ExtensionAPI) {
-  let phase: WorkflowPhase = 'off';
-  let restoreToolNames: string[] | null = null;
-  let todoItems: TodoItem[] = [];
-  let returnToPlanningAfterNextAgentEnd = false;
-
-  function getPlanCommandCompletions(argumentText: string): AutocompleteItem[] | null {
-    if (argumentText.trim().includes(' ')) return null;
-
-    const query = argumentText.trim().toLowerCase();
-    const items: AutocompleteItem[] = [
-      {
-        value: 'status',
-        label: 'status',
-        description: `Show the current plan workflow state (currently ${phase})`,
-      },
-      {
-        value: 'domain',
-        label: 'domain',
-        description: 'Stress-test the current plan with a domain-model review',
-      },
-      {
-        value: 'plans',
-        label: 'plans',
-        description: 'Generate self-contained implementation plan files',
-      },
-      {
-        value: 'execute',
-        label: 'execute',
-        description: 'Leave read-only planning and execute the approved plan',
-      },
-      {
-        value: 'off',
-        label: 'off',
-        description: 'Disable plan mode',
-      },
-    ];
-
-    const filtered = items.filter((item) => {
-      if (!query) return true;
-      return (
-        item.value.toLowerCase().startsWith(query) ||
-        item.label?.toLowerCase().includes(query) ||
-        item.description?.toLowerCase().includes(query)
-      );
-    });
-
-    return filtered.length > 0 ? filtered : null;
-  }
+// ── Extension ────────────────────────────────────────────────────────────────
+export default function planMode(pi: ExtensionAPI): void {
+  let planEnabled = false;
+  let executing = false;
+  let planDir: string | undefined;
+  let todos: TodoItem[] = [];
+  let previousThinking: ThinkingLevel | undefined;
+  let previousModel: { provider: string; id: string } | undefined;
 
+  // ── Flag ──────────────────────────────────────────────────────────────────
   pi.registerFlag('plan', {
-    description: 'Start in plan mode (Cursor-style planning workflow)',
+    description: 'Start in plan mode (read-only + medium thinking)',
     type: 'boolean',
     default: false,
   });
 
-  function persistState() {
-    pi.appendEntry<PersistedState>(STATE_ENTRY, {
-      phase,
-      todos: todoItems.map((item) => ({ ...item })),
+  // ── State persistence ─────────────────────────────────────────────────────
+  function persist(): void {
+    pi.appendEntry<PersistedState>('plan-mode', {
+      planEnabled,
+      executing,
+      planDir,
+      todos,
     });
   }
 
-  function updateUi(ctx: ExtensionContext) {
-    if (phase === 'off') {
-      ctx.ui.setStatus('plan-mode', undefined);
-      ctx.ui.setWidget('plan-mode-todos', undefined);
-      return;
-    }
+  // ── UI updates ────────────────────────────────────────────────────────────
+  function updateUI(ctx: ExtensionContext): void {
+    const { theme } = ctx.ui;
 
-    if (phase === 'planning') {
-      ctx.ui.setStatus('plan-mode', ctx.ui.theme.fg('warning', 'plan'));
-    } else if (phase === 'plan-files') {
-      ctx.ui.setStatus('plan-mode', ctx.ui.theme.fg('accent', 'plan:files'));
+    if (executing && todos.length > 0) {
+      const done = todos.filter((t) => t.completed).length;
+      ctx.ui.setStatus('plan-mode', theme.fg('accent', `📋 exec ${done}/${todos.length}`));
+    } else if (planEnabled) {
+      ctx.ui.setStatus('plan-mode', theme.fg('warning', '📝 plan'));
     } else {
-      const completed = todoItems.filter((item) => item.completed).length;
-      const total = todoItems.length;
-      ctx.ui.setStatus(
-        'plan-mode',
-        ctx.ui.theme.fg('success', total > 0 ? `plan:exec ${completed}/${total}` : 'plan:exec'),
-      );
-    }
-
-    if (todoItems.length === 0) {
-      ctx.ui.setWidget('plan-mode-todos', undefined);
-      return;
-    }
-
-    const lines = todoItems.map((item) => {
-      if (item.completed) {
-        return ctx.ui.theme.fg('success', `☑ ${ctx.ui.theme.strikethrough(item.text)}`);
-      }
-      const marker = phase === 'executing' ? '☐' : '•';
-      return `${ctx.ui.theme.fg('muted', `${marker} `)}${item.text}`;
-    });
-    ctx.ui.setWidget('plan-mode-todos', lines);
-  }
-
-  function ensureRestoreTools() {
-    if (restoreToolNames) return;
-    restoreToolNames = [...pi.getActiveTools()];
-  }
-
-  function applyPhaseTools(targetPhase: WorkflowPhase) {
-    const allToolNames = pi.getAllTools().map((tool) => tool.name);
-
-    if (targetPhase === 'planning') {
-      const resolved = resolveToolNames(allToolNames, [...PLANNING_TOOL_PATTERNS, 'context7_*']);
-      if (resolved.length > 0) pi.setActiveTools(resolved);
-      return;
-    }
-
-    if (targetPhase === 'plan-files') {
-      const resolved = resolveToolNames(allToolNames, [...PLAN_FILE_TOOL_PATTERNS, 'context7_*']);
-      if (resolved.length > 0) pi.setActiveTools(resolved);
-      return;
+      ctx.ui.setStatus('plan-mode', undefined);
     }
 
-    if (targetPhase === 'executing' && restoreToolNames) {
-      pi.setActiveTools([...restoreToolNames]);
+    if (executing && todos.length > 0) {
+      const lines = todos.map((item) => {
+        if (item.completed) {
+          return theme.fg('success', '☑ ') + theme.fg('muted', theme.strikethrough(item.text));
+        }
+        return `${theme.fg('muted', '☐ ')}${item.text}`;
+      });
+      ctx.ui.setWidget('plan-todos', lines);
+    } else {
+      ctx.ui.setWidget('plan-todos', undefined);
     }
   }
 
-  function setPhase(
-    nextPhase: WorkflowPhase,
+  // ── Model switching ───────────────────────────────────────────────────────
+  async function switchModel(
     ctx: ExtensionContext,
-    options?: { notify?: string },
-  ) {
-    if (phase === 'off' && nextPhase !== 'off') ensureRestoreTools();
-    if (nextPhase === 'off' && restoreToolNames) {
-      pi.setActiveTools([...restoreToolNames]);
-      restoreToolNames = null;
-    } else {
-      applyPhaseTools(nextPhase);
+    preset: { provider: string; id: string },
+  ): Promise<boolean> {
+    const model = ctx.modelRegistry.find(preset.provider, preset.id);
+    if (!model) {
+      ctx.ui.notify(`Model ${preset.provider}/${preset.id} not found`, 'error');
+      return false;
     }
-
-    phase = nextPhase;
-    updateUi(ctx);
-    persistState();
-
-    if (options?.notify) ctx.ui.notify(options.notify, 'info');
-  }
-
-  async function sendPlanningPrompt(prompt: string, ctx: ExtensionContext) {
-    const commandCtx = ctx as ExtensionContext & { waitForIdle?: () => Promise<void> };
-    if (typeof commandCtx.waitForIdle === 'function') {
-      await commandCtx.waitForIdle();
+    const ok = await pi.setModel(model);
+    if (!ok) {
+      ctx.ui.notify(`No API key for ${preset.provider}/${preset.id}`, 'error');
+      return false;
     }
-    pi.sendUserMessage(prompt.trim());
-  }
-
-  function availableCommand(name: string): string | undefined {
-    const commands = pi.getCommands?.() ?? [];
-    return commands.find((command) => command.name === name)?.name;
-  }
-
-  async function startPlanning(prompt: string | undefined, ctx: ExtensionContext) {
-    setPhase('planning', ctx, {
-      notify:
-        phase === 'off'
-          ? 'Plan mode enabled. Ask pi to inspect, question, and plan before implementation.'
-          : undefined,
-    });
-
-    const trimmed = prompt?.trim();
-    if (!trimmed) return;
-    await sendPlanningPrompt(trimmed, ctx);
+    return true;
   }
 
-  async function runDomainWorkflow(args: string | undefined, ctx: ExtensionContext) {
-    setPhase('planning', ctx, { notify: 'Running domain-model review in read-only plan mode.' });
-
-    const planText =
-      todoItems.length > 0 ? formatTodoList(todoItems) : 'No extracted plan steps yet.';
-    const trimmedArgs = args?.trim();
-    const promptText =
-      trimmedArgs && trimmedArgs.length > 0
-        ? trimmedArgs
-        : `Stress-test the current plan against the existing domain model and terminology. Challenge ambiguous terms, invent concrete scenarios, compare claims against the codebase, and suggest any CONTEXT.md or ADR updates only when justified.\n\nCurrent plan:\n${planText}`;
-    const skillCommand = availableCommand('skill:domain-model');
-
-    if (skillCommand) {
-      const suffix =
-        trimmedArgs && trimmedArgs.length > 0
-          ? trimmedArgs
-          : `Stress-test the current plan against the domain model.\n\nCurrent plan:\n${planText}`;
-      await sendPlanningPrompt(`/${skillCommand} ${suffix}`.trim(), ctx);
-      return;
-    }
-
-    await sendPlanningPrompt(promptText, ctx);
+  // ── Phase transitions ─────────────────────────────────────────────────────
+  async function enterPlanMode(ctx: ExtensionContext): Promise<void> {
+    planEnabled = true;
+    executing = false;
+    planDir = undefined;
+    todos = [];
+    previousThinking = pi.getThinkingLevel() as ThinkingLevel;
+    previousModel = ctx.model ? { provider: ctx.model.provider, id: ctx.model.id } : undefined;
+    pi.setActiveTools(PLAN_TOOLS);
+    await switchModel(ctx, PLAN_MODEL);
+    pi.setThinkingLevel(PLAN_THINKING);
+    ctx.ui.notify(
+      `Plan mode ON — ${PLAN_MODEL.provider}/${PLAN_MODEL.id}:${PLAN_THINKING}`,
+      'info',
+    );
+    updateUI(ctx);
+    persist();
   }
 
-  async function runPlanFileWorkflow(args: string | undefined, ctx: ExtensionContext) {
-    setPhase('plan-files', ctx, {
-      notify:
-        'Plan-file authoring enabled. pi may write planning docs, but should not implement code.',
-    });
-    returnToPlanningAfterNextAgentEnd = true;
-
-    const planText =
-      todoItems.length > 0 ? formatTodoList(todoItems) : 'No extracted plan steps yet.';
-    const trimmedArgs = args?.trim();
-    const promptText =
-      trimmedArgs && trimmedArgs.length > 0
-        ? trimmedArgs
-        : `Create self-contained implementation plan files for the current approved plan. Ground every file in the real codebase, document exact APIs and file paths, and write the plans to docs/plans unless the repo or user prefers another location.\n\nCurrent plan:\n${planText}`;
-    const skillCommand = availableCommand('skill:create-implementation-plans');
-
-    if (skillCommand) {
-      const suffix =
-        trimmedArgs && trimmedArgs.length > 0
-          ? trimmedArgs
-          : `Create self-contained implementation plan files for the current approved plan.\n\nCurrent plan:\n${planText}`;
-      await sendPlanningPrompt(`/${skillCommand} ${suffix}`.trim(), ctx);
-      return;
+  async function exitPlanMode(ctx: ExtensionContext): Promise<void> {
+    planEnabled = false;
+    executing = false;
+    planDir = undefined;
+    todos = [];
+    pi.setActiveTools(EXEC_TOOLS);
+    if (previousModel) {
+      await switchModel(ctx, previousModel);
     }
-
-    await sendPlanningPrompt(promptText, ctx);
-  }
-
-  async function startExecution(args: string | undefined, ctx: ExtensionContext) {
-    setPhase('executing', ctx, {
-      notify: 'Plan execution mode enabled. Full tool access restored.',
-    });
-
-    const trimmedArgs = args?.trim();
-    if (trimmedArgs && trimmedArgs.length > 0) {
-      await sendPlanningPrompt(trimmedArgs, ctx);
-      return;
+    if (previousThinking) {
+      pi.setThinkingLevel(previousThinking);
     }
-
-    const defaultPrompt =
-      todoItems.length > 0
-        ? `Execute the approved plan in order.\n\nPlan:\n${formatTodoList(todoItems)}\n\nStart with step 1 and include [DONE:n] tags as each step is completed.`
-        : 'Execute the approved plan in small verified steps. Include [DONE:n] tags as each major plan step is completed.';
-    await sendPlanningPrompt(defaultPrompt, ctx);
+    ctx.ui.notify('Plan mode OFF — original model restored', 'info');
+    updateUI(ctx);
+    persist();
   }
 
-  function showStatus(ctx: ExtensionContext) {
-    const lines = [`Phase: ${phase}`];
-    if (todoItems.length > 0) {
-      const completed = todoItems.filter((item) => item.completed).length;
-      lines.push(`Plan steps: ${completed}/${todoItems.length} complete`);
-      lines.push(
-        '',
-        ...todoItems.map((item) => `${item.step}. ${item.completed ? '✓' : '○'} ${item.text}`),
-      );
-    }
-    ctx.ui.notify(lines.join('\n'), 'info');
+  async function startExecution(ctx: ExtensionContext): Promise<void> {
+    planEnabled = false;
+    executing = true;
+    pi.setActiveTools(EXEC_TOOLS);
+    await switchModel(ctx, EXEC_MODEL);
+    pi.setThinkingLevel(EXEC_THINKING);
+    ctx.ui.notify(
+      `Executing plan — ${EXEC_MODEL.provider}/${EXEC_MODEL.id}:${EXEC_THINKING}`,
+      'info',
+    );
+    updateUI(ctx);
+    persist();
   }
 
-  async function showPlanMenu(ctx: ExtensionContext) {
-    if (!ctx.hasUI) {
-      showStatus(ctx);
-      return;
-    }
-
-    const choice = await ctx.ui.select('Plan mode', [
-      'Stay in planning mode',
-      'Stress-test with domain-model',
-      'Generate implementation plan files',
-      'Execute current plan',
-      'Show status',
-      'Disable plan mode',
-    ]);
-
-    if (!choice || choice === 'Stay in planning mode') return;
-    if (choice === 'Stress-test with domain-model') {
-      await runDomainWorkflow(undefined, ctx);
-      return;
-    }
-    if (choice === 'Generate implementation plan files') {
-      await runPlanFileWorkflow(undefined, ctx);
-      return;
-    }
-    if (choice === 'Execute current plan') {
-      await startExecution(undefined, ctx);
-      return;
-    }
-    if (choice === 'Show status') {
-      showStatus(ctx);
-      return;
+  async function togglePlanMode(ctx: ExtensionContext): Promise<void> {
+    if (planEnabled || executing) {
+      await exitPlanMode(ctx);
+    } else {
+      await enterPlanMode(ctx);
     }
-
-    setPhase('off', ctx, { notify: 'Plan mode disabled.' });
   }
 
+  // ── Commands ──────────────────────────────────────────────────────────────
   pi.registerCommand('plan', {
-    description: 'Enable plan mode, send a planning prompt, or manage the plan workflow',
-    getArgumentCompletions: getPlanCommandCompletions,
+    description: 'Enter plan mode, optionally with a starting prompt',
     handler: async (args, ctx) => {
-      const raw = args?.trim() ?? '';
-      if (!raw) {
-        if (phase === 'off') {
-          await startPlanning(undefined, ctx);
-        } else {
-          await showPlanMenu(ctx);
-        }
+      if (planEnabled || executing) {
+        await togglePlanMode(ctx);
         return;
       }
-
-      const lower = raw.toLowerCase();
-      if (CLEAR_VALUES.has(lower)) {
-        setPhase('off', ctx, { notify: 'Plan mode disabled.' });
-        return;
+      await enterPlanMode(ctx);
+      const prompt = args?.trim();
+      if (prompt) {
+        pi.sendUserMessage(prompt);
       }
-      if (lower === 'status') {
-        showStatus(ctx);
-        return;
-      }
-      if (lower === 'domain') {
-        await runDomainWorkflow(undefined, ctx);
-        return;
-      }
-      if (lower === 'plans' || lower === 'files') {
-        await runPlanFileWorkflow(undefined, ctx);
-        return;
-      }
-      if (lower === 'execute' || lower === 'run') {
-        await startExecution(undefined, ctx);
-        return;
-      }
-
-      await startPlanning(raw, ctx);
     },
   });
 
-  pi.registerCommand('plan-status', {
-    description: 'Show current plan workflow phase and extracted plan steps',
+  pi.registerCommand('todos', {
+    description: 'Show current plan progress',
     handler: async (_args, ctx) => {
-      showStatus(ctx);
-    },
-  });
-
-  pi.registerCommand('plan-domain', {
-    description: 'Run a domain-model stress test for the current plan',
-    handler: async (args, ctx) => {
-      await runDomainWorkflow(args, ctx);
-    },
-  });
-
-  pi.registerCommand('plan-plans', {
-    description: 'Generate self-contained implementation plan files for the current plan',
-    handler: async (args, ctx) => {
-      await runPlanFileWorkflow(args, ctx);
-    },
-  });
-
-  pi.registerCommand('plan-execute', {
-    description: 'Leave read-only planning and execute the current approved plan',
-    handler: async (args, ctx) => {
-      await startExecution(args, ctx);
+      if (todos.length === 0) {
+        ctx.ui.notify('No plan yet. Use /plan to start planning.', 'info');
+        return;
+      }
+      const list = todos
+        .map((t, i) => `${i + 1}. ${t.completed ? '✓' : '○'} ${t.text}`)
+        .join('\n');
+      ctx.ui.notify(`Plan Progress:\n${list}`, 'info');
     },
   });
 
   pi.registerShortcut(Key.ctrlAlt('p'), {
     description: 'Toggle plan mode',
-    handler: async (ctx) => {
-      if (phase === 'off') setPhase('planning', ctx, { notify: 'Plan mode enabled.' });
-      else setPhase('off', ctx, { notify: 'Plan mode disabled.' });
-    },
+    handler: async (ctx) => togglePlanMode(ctx),
   });
 
+  // ── Block destructive bash in plan mode ───────────────────────────────────
   pi.on('tool_call', async (event) => {
-    if (phase !== 'planning' || event.toolName !== 'bash') return;
+    if (!planEnabled) return;
+
+    // Block bash commands that aren't on the safe allowlist
+    if (event.toolName === 'bash') {
+      const command = event.input.command as string;
+      if (!isSafeCommand(command)) {
+        return {
+          block: true,
+          reason: `Plan mode: command blocked. Use /plan to exit plan mode first.\nCommand: ${command}`,
+        };
+      }
+    }
 
-    const command = event.input.command as string;
-    if (isSafeCommand(command)) return;
+    // Block edit/write to paths outside .plans/
+    if (event.toolName === 'edit' || event.toolName === 'write') {
+      const path = (event.input as { path?: string }).path ?? '';
+      if (!path.startsWith('.plans/') && !path.startsWith('.plans\\')) {
+        return {
+          block: true,
+          reason: `Plan mode: file modifications are restricted to .plans/ directory.\nPath: ${path}`,
+        };
+      }
+    }
+  });
 
+  // ── Filter stale plan context when not planning ───────────────────────────
+  pi.on('context', async (event) => {
+    if (planEnabled) return;
     return {
-      block: true,
-      reason: `Plan mode blocks non-read-only bash commands. Disable plan mode or use /plan-plans for controlled plan-file authoring.\nCommand: ${command}`,
+      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;
+      }),
     };
   });
 
-  pi.on('before_agent_start', async (event, ctx) => {
-    if (phase === 'off') return;
+  // ── Inject context for each phase ─────────────────────────────────────────
+  pi.on('before_agent_start', async () => {
+    if (planEnabled) {
+      return {
+        message: {
+          customType: 'plan-mode-context',
+          content: `[PLAN MODE ACTIVE]
+You are in plan mode — a planning phase with strict bash restrictions.
 
-    const compatEvent = event as typeof event & BeforeAgentStartCompatEvent;
-    const activeTools = compatEvent.systemPromptOptions?.selectedTools ?? pi.getActiveTools();
+Restrictions:
+- Available tools: ${PLAN_TOOLS.join(', ')}
+- Bash is restricted to read-only commands (ls, grep, git status, etc.)
+- edit and write are ONLY allowed for files inside the \`.plans/\` directory
 
-    if (phase === 'planning') {
-      return {
-        systemPrompt: `${event.systemPrompt}\n\n${buildPlanningInstructions(ctx, activeTools, todoItems)}`,
+Your task:
+1. Analyze the codebase thoroughly using the available read-only tools
+2. Ask clarifying questions if needed (use the questionnaire tool)
+3. Produce a detailed, concrete plan
+
+When you are ready to finalize the plan:
+1. Choose a short descriptive kebab-case name for the plan (e.g. "add-auth-middleware")
+2. Create \`.plans/<plan-name>/PLAN.md\` with the full numbered plan under a \`Plan:\` header:
+
+\`\`\`markdown
+# <Plan Title>
+
+<Brief description of what this plan accomplishes>
+
+## Context
+<Key findings from codebase analysis>
+
+## Plan:
+1. First step — what to change and where
+2. Second step — what to change and where
+...
+
+## Risks / Open Questions
+<Any concerns or assumptions>
+\`\`\`
+
+3. Create \`.plans/<plan-name>/START-PROMPT.md\` — a self-contained handoff prompt that a different model can use to execute the plan WITHOUT access to this conversation. It must include:
+   - Complete context about the codebase (relevant file paths, APIs, patterns)
+   - The full plan steps to execute
+   - Any critical constraints or gotchas
+   - Clear instructions to mark each step done with \`[DONE:n]\` tags
+
+The START-PROMPT.md is critical — it must be thorough enough that an implementor with zero prior context can execute the plan correctly.
+
+If you need supporting reference files for extra context (code snippets, diagrams, specs), place them alongside in the same \`.plans/<plan-name>/\` directory.
+
+Do NOT attempt to make product code changes — only create planning artifacts in \`.plans/\`.`,
+          display: false,
+        },
       };
     }
 
-    if (phase === 'plan-files') {
+    if (executing && todos.length > 0) {
+      const remaining = todos.filter((t) => !t.completed);
+      const todoList = remaining.map((t) => `${t.step}. ${t.text}`).join('\n');
       return {
-        systemPrompt: `${event.systemPrompt}\n\n${buildPlanFileInstructions(activeTools, todoItems)}`,
+        message: {
+          customType: 'plan-execution-context',
+          content: `[EXECUTING PLAN — Full tool access enabled]
+
+Remaining steps:
+${todoList}
+
+Execute each step in order. You MUST include [DONE:n] in your response after completing each step before moving to the next one.`,
+          display: false,
+        },
       };
     }
-
-    return {
-      systemPrompt: `${event.systemPrompt}\n\n${buildExecutionInstructions(todoItems)}`,
-    };
   });
 
+  // ── Track [DONE:n] markers during execution ───────────────────────────────
   pi.on('turn_end', async (event, ctx) => {
-    if (phase !== 'executing' || todoItems.length === 0) return;
-    const text = getMessageText(event.message);
-    if (!text) return;
+    if (!executing || todos.length === 0) return;
+    if (!isAssistantMessage(event.message)) return;
 
-    if (markCompletedSteps(text, todoItems) > 0) {
-      updateUi(ctx);
-      persistState();
+    const text = getTextContent(event.message);
+    if (markCompletedSteps(text, todos) > 0) {
+      updateUI(ctx);
     }
+    persist();
   });
 
-  pi.on('agent_end', async (event, ctx) => {
-    if (returnToPlanningAfterNextAgentEnd && phase === 'plan-files') {
-      returnToPlanningAfterNextAgentEnd = false;
-      setPhase('planning', ctx, { notify: 'Returned to read-only plan mode.' });
+  // ── Detect plan directory from written files ──────────────────────────────
+  pi.on('tool_result', async (event) => {
+    if (!planEnabled) return;
+    if (event.toolName !== 'write' && event.toolName !== 'edit') return;
+    if (event.isError) return;
+
+    const path = (event.input as { path?: string }).path;
+    if (!path) return;
+
+    // Detect .plans/<name>/ directory from written files
+    const match = path.match(/\.plans\/([^/]+)\//);
+    if (match && !planDir) {
+      planDir = `.plans/${match[1]}`;
+      persist();
     }
+  });
 
-    if (
-      phase === 'executing' &&
-      todoItems.length > 0 &&
-      todoItems.every((item) => item.completed)
-    ) {
-      setPhase('off', ctx, { notify: 'Plan execution complete.' });
-      todoItems = [];
-      updateUi(ctx);
-      persistState();
+  // ── After agent finishes: prompt for next action ──────────────────────────
+  pi.on('agent_end', async (event, ctx) => {
+    // Check execution completion
+    if (executing && todos.length > 0) {
+      if (todos.every((t) => t.completed)) {
+        const list = todos.map((t) => `~~${t.text}~~`).join('\n');
+        pi.sendMessage(
+          {
+            customType: 'plan-complete',
+            content: `**Plan Complete!** ✓\n\n${list}`,
+            display: true,
+          },
+          { triggerTurn: false },
+        );
+        executing = false;
+        todos = [];
+        planDir = undefined;
+        pi.setActiveTools(EXEC_TOOLS);
+        if (previousModel) {
+          await switchModel(ctx, previousModel);
+        }
+        if (previousThinking) {
+          pi.setThinkingLevel(previousThinking);
+        }
+        updateUI(ctx);
+        persist();
+      }
       return;
     }
 
-    if (phase !== 'planning') return;
-
-    const lastAssistantMessage = [...event.messages]
-      .reverse()
-      .find((message) => isAssistantMessage(message));
-    if (!lastAssistantMessage) return;
-
-    const extracted = extractTodoItems(getMessageText(lastAssistantMessage));
-    if (extracted.length > 0) {
-      todoItems = extracted;
-      updateUi(ctx);
-      persistState();
-    }
+    if (!planEnabled || !ctx.hasUI) return;
 
-    if (todoItems.length === 0 || !ctx.hasUI) return;
+    // Check if plan files were created by looking for planDir
+    if (!planDir) return;
 
+    // Show menu
     const choice = await ctx.ui.select('Plan ready — what next?', [
-      'Stay in planning mode',
-      'Stress-test with domain-model',
-      'Generate implementation plan files',
-      'Execute current plan',
-      'Refine the plan',
-      'Disable plan mode',
+      'Execute Plan',
+      'Refine Plan',
+      'Follow up',
+      'Exit plan mode',
     ]);
 
-    if (!choice || choice === 'Stay in planning mode') return;
-    if (choice === 'Stress-test with domain-model') {
-      await runDomainWorkflow(undefined, ctx);
-      return;
-    }
-    if (choice === 'Generate implementation plan files') {
-      await runPlanFileWorkflow(undefined, ctx);
-      return;
-    }
-    if (choice === 'Execute current plan') {
-      await startExecution(undefined, ctx);
-      return;
-    }
-    if (choice === 'Refine the plan') {
-      const refinement = await ctx.ui.editor('Refine the plan:', '');
-      if (refinement?.trim()) await sendPlanningPrompt(refinement, ctx);
-      return;
-    }
+    if (choice === 'Execute Plan') {
+      // Read START-PROMPT.md for clean context handoff
+      const startPromptPath = `${planDir}/START-PROMPT.md`;
+      const planMdPath = `${planDir}/PLAN.md`;
+
+      // Read the plan to extract todos
+      let planContent = '';
+      try {
+        const result = await pi.exec('cat', [planMdPath]);
+        if (result.code === 0) {
+          planContent = result.stdout;
+        }
+      } catch {
+        // Fall through — will use empty plan content
+      }
 
-    setPhase('off', ctx, { notify: 'Plan mode disabled.' });
-  });
+      const extracted = extractTodoItems(planContent);
+      if (extracted.length > 0) {
+        todos = extracted;
+      }
 
-  async function restoreState(ctx: ExtensionContext) {
-    const savedState = findSavedState(ctx);
+      // Read the start prompt for clean handoff
+      let startPrompt = '';
+      try {
+        const result = await pi.exec('cat', [startPromptPath]);
+        if (result.code === 0) {
+          startPrompt = result.stdout.trim();
+        }
+      } catch {
+        // Fall through
+      }
 
-    if (pi.getFlag('plan') === true && !savedState) {
-      ensureRestoreTools();
-      phase = 'planning';
-      applyPhaseTools(phase);
-      updateUi(ctx);
-      persistState();
-      return;
+      await startExecution(ctx);
+      updateUI(ctx);
+
+      if (startPrompt) {
+        pi.sendMessage(
+          {
+            customType: 'plan-mode-execute',
+            content: startPrompt,
+            display: true,
+          },
+          { triggerTurn: true },
+        );
+      } else {
+        // Fallback: ask executor to read the plan
+        pi.sendMessage(
+          {
+            customType: 'plan-mode-execute',
+            content: `Execute the plan in ${planMdPath}. Read it first, then execute step by step. Mark each step with [DONE:n] before moving to the next.`,
+            display: true,
+          },
+          { triggerTurn: true },
+        );
+      }
+    } else if (choice === 'Refine Plan') {
+      // Adversarial review — planner critiques its own plan
+      pi.sendMessage(
+        {
+          customType: 'plan-mode-refine',
+          content: `Review the plan you just created in ${planDir}/PLAN.md with an adversarial lens. Challenge assumptions, find gaps, identify risks, and look for:
+
+- Missing edge cases or error handling
+- Incorrect assumptions about the codebase
+- Steps that are too vague or could be misinterpreted
+- Missing dependencies between steps
+- Simpler alternatives that were overlooked
+
+After your review, update PLAN.md and START-PROMPT.md with any improvements.`,
+          display: true,
+        },
+        { triggerTurn: true },
+      );
+    } else if (choice === 'Follow up') {
+      const followUp = await ctx.ui.editor('Follow-up instructions for the planner:', '');
+      if (followUp?.trim()) {
+        pi.sendUserMessage(followUp.trim());
+      }
+    } else if (choice === 'Exit plan mode') {
+      await exitPlanMode(ctx);
     }
+  });
 
-    if (!savedState) {
-      phase = 'off';
-      todoItems = [];
-      updateUi(ctx);
-      return;
+  // ── Restore state on session start/resume ─────────────────────────────────
+  pi.on('session_start', async (_event, ctx) => {
+    // Check CLI flag
+    if (pi.getFlag('plan') === true) {
+      planEnabled = true;
     }
 
-    todoItems = savedState.todos.map((item) => ({ ...item }));
-    phase = savedState.phase;
-
-    if (phase === 'planning' || phase === 'plan-files') {
-      ensureRestoreTools();
-      applyPhaseTools(phase);
+    // Restore persisted state
+    const entries = ctx.sessionManager.getEntries();
+    const saved = entries
+      .filter(
+        (e: { type: string; customType?: string }) =>
+          e.type === 'custom' && e.customType === 'plan-mode',
+      )
+      .pop() as { data?: PersistedState } | undefined;
+
+    if (saved?.data) {
+      planEnabled = saved.data.planEnabled ?? planEnabled;
+      executing = saved.data.executing ?? executing;
+      planDir = saved.data.planDir ?? planDir;
+      todos = saved.data.todos ?? todos;
     }
 
-    updateUi(ctx);
-  }
+    // Re-scan [DONE:n] markers on resume
+    if (executing && todos.length > 0) {
+      let execIdx = -1;
+      for (let i = entries.length - 1; i >= 0; i--) {
+        const entry = entries[i] as { type: string; customType?: string };
+        if (entry.customType === 'plan-mode-execute') {
+          execIdx = i;
+          break;
+        }
+      }
 
-  pi.on('session_start', async (_event, ctx) => {
-    await restoreState(ctx);
-  });
+      const messages: AssistantMessage[] = [];
+      for (let i = execIdx + 1; i < entries.length; i++) {
+        const entry = entries[i];
+        if (
+          entry.type === 'message' &&
+          'message' in entry &&
+          isAssistantMessage(entry.message as AgentMessage)
+        ) {
+          messages.push(entry.message as AssistantMessage);
+        }
+      }
+      const allText = messages.map(getTextContent).join('\n');
+      markCompletedSteps(allText, todos);
+    }
+
+    // Apply tool restrictions, model, and thinking level
+    if (planEnabled) {
+      pi.setActiveTools(PLAN_TOOLS);
+      await switchModel(ctx, PLAN_MODEL);
+      pi.setThinkingLevel(PLAN_THINKING);
+    } else if (executing) {
+      pi.setActiveTools(EXEC_TOOLS);
+      await switchModel(ctx, EXEC_MODEL);
+      pi.setThinkingLevel(EXEC_THINKING);
+    }
 
-  pi.on('session_tree', async (_event, ctx) => {
-    await restoreState(ctx);
+    updateUI(ctx);
   });
 }