@docyrus/docyrus 0.0.25 → 0.0.26

package/resources/pi-agent/extensions/plan.ts

@@ -0,0 +1,1197 @@
+import type {
+  AgentEndEvent,
+  ExtensionAPI,
+  ExtensionCommandContext,
+  ExtensionContext,
+  ThinkingLevel,
+  ToolCallEvent,
+  UserBashEvent,
+} from "@mariozechner/pi-coding-agent";
+import { DynamicBorder } from "@mariozechner/pi-coding-agent";
+import type { Api, Model } from "@mariozechner/pi-ai";
+import { Container, Text } from "@mariozechner/pi-tui";
+import { existsSync, readFileSync } from "node:fs";
+import fs from "node:fs/promises";
+import os from "node:os";
+import path from "node:path";
+import {
+  buildAskUserProtocolInstructions,
+  formatAskUserResponsePrompt,
+  parseAskUserRequestFromText,
+  type IAskUserQuestion,
+  type IAskUserResponse,
+} from "../shared/askUserProtocol";
+
+const PLAN_OUTPUT_ROOT_SEGMENTS = [".docyrus", "plans"] as const;
+const PLAN_STATE_TYPE = "plan-session";
+const PLAN_ANCHOR_TYPE = "plan-anchor";
+const PLAN_CONFIG_FILE = ".pi/plan-policy.json";
+const PLAN_WIDGET_KEY = "plan-mode";
+const PLAN_BRANCH_LABEL = "plan";
+const ALLOWED_TODO_ACTIONS = new Set(["list", "list-all", "get"]);
+const ALL_THINKING_LEVELS = new Set(["off", "minimal", "low", "medium", "high", "xhigh"] satisfies ThinkingLevel[]);
+
+const PLAN_MODE_SYSTEM_PROMPT = [
+  "You are in Docyrus `/plan` mode.",
+  "",
+  "This is planning only.",
+  "Do not implement code, do not mutate repository files, and do not perform remote Docyrus mutations.",
+  "Use inspection, analysis, and verification commands only.",
+  "",
+  "Respond with an actionable Markdown plan using these exact top-level sections in this order:",
+  "## Summary",
+  "## Implementation Changes",
+  "## Test Plan",
+  "## Assumptions",
+  "",
+  "Keep the plan decision-complete for implementation, call out assumptions explicitly, and do not start implementation.",
+].join("\n");
+
+const END_PLAN_SUMMARY_PROMPT = `You are returning from a planning-only branch.
+
+Summarize the plan branch as a concise, implementation-ready plan.
+
+Requirements:
+- Do not implement anything.
+- Preserve the important decisions made on the plan branch.
+- Output Markdown with these exact top-level sections:
+  - Summary
+  - Implementation Changes
+  - Test Plan
+  - Assumptions
+- Keep it concise and actionable.`;
+
+export interface IPlanModelEntry {
+	model: string;
+	thinkingLevel?: ThinkingLevel;
+}
+
+export interface IPlanProfileOverride {
+	match: string;
+	models?: IPlanModelEntry[];
+}
+
+export interface IPlanPolicy {
+	models?: IPlanModelEntry[];
+	profiles?: Record<string, IPlanProfileOverride>;
+}
+
+export interface IPlanSessionState {
+	active: boolean;
+	mode?: "plan" | "architect";
+	originId?: string;
+	branchStartId?: string;
+	artifactPath?: string;
+	task?: string;
+	sourceModel?: string;
+	sourceThinking?: ThinkingLevel;
+	planModel?: string;
+	planThinking?: ThinkingLevel;
+	startedAt?: string;
+}
+
+export interface IReadPlanPolicyResult {
+	policy: IPlanPolicy;
+	path?: string;
+	error?: string;
+}
+
+export type IParseResult<T> = { ok: true; value: T } | { ok: false; error: string };
+
+let currentPlanState: IPlanSessionState | undefined;
+
+export interface IPlanningStartParams {
+	pi: ExtensionAPI;
+	ctx: ExtensionCommandContext;
+	mode: "plan" | "architect";
+	task?: string;
+	artifactPath: string;
+	initialPrompt: string;
+	profileName?: string;
+}
+
+function expandUserPath(inputPath: string): string {
+  if (inputPath === "~") {
+    return os.homedir();
+  }
+  if (inputPath.startsWith("~/")) {
+    return path.join(os.homedir(), inputPath.slice(2));
+  }
+  return inputPath;
+}
+
+function resolveAgentRootPath(): string {
+  const agentDir = process.env.PI_CODING_AGENT_DIR?.trim();
+  return agentDir && agentDir.length > 0 ? expandUserPath(agentDir) : path.join(os.homedir(), ".pi", "agent");
+}
+
+function resolveGlobalPlanPolicyPath(): string {
+  return path.join(resolveAgentRootPath(), "plan-policy.json");
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+export function parsePlanTask(rawArgs: string): string | null {
+  const normalized = rawArgs.trim();
+  return normalized.length > 0 ? normalized : null;
+}
+
+export function slugifyPlanTask(task: string): string {
+  const slug = task
+    .toLowerCase()
+    .replace(/[^a-z0-9]+/g, "-")
+    .replace(/^-+|-+$/g, "")
+    .slice(0, 60);
+
+  return slug || "plan";
+}
+
+export function formatPlanTimestamp(date: Date): string {
+  const year = `${date.getFullYear()}`;
+  const month = `${date.getMonth() + 1}`.padStart(2, "0");
+  const day = `${date.getDate()}`.padStart(2, "0");
+  const hour = `${date.getHours()}`.padStart(2, "0");
+  const minute = `${date.getMinutes()}`.padStart(2, "0");
+  const second = `${date.getSeconds()}`.padStart(2, "0");
+  return `${year}${month}${day}-${hour}${minute}${second}`;
+}
+
+export function createPlanArtifactPath(params: {
+	cwd: string;
+	task: string;
+	date?: Date;
+}): string {
+  const timestamp = formatPlanTimestamp(params.date ?? new Date());
+  const slug = slugifyPlanTask(params.task);
+  return path.join(params.cwd, ...PLAN_OUTPUT_ROOT_SEGMENTS, `${timestamp}-${slug}.md`);
+}
+
+export function parsePlanModelSelector(value: unknown): IParseResult<string> {
+  if (typeof value !== "string") {
+    return { ok: false, error: "expected model selector provider/modelId" };
+  }
+  if (value.trim() !== value) {
+    return { ok: false, error: "expected model selector provider/modelId" };
+  }
+  const slashIndex = value.indexOf("/");
+  if (slashIndex <= 0 || slashIndex >= value.length - 1) {
+    return { ok: false, error: "expected model selector provider/modelId" };
+  }
+  const provider = value.slice(0, slashIndex);
+  const modelId = value.slice(slashIndex + 1);
+  if (/\s/.test(provider) || /\s/.test(modelId)) {
+    return { ok: false, error: "expected model selector provider/modelId" };
+  }
+  return { ok: true, value };
+}
+
+function parseThinkingLevel(value: unknown): IParseResult<ThinkingLevel> {
+  if (typeof value !== "string" || !ALL_THINKING_LEVELS.has(value as ThinkingLevel)) {
+    return { ok: false, error: "expected one of: off, minimal, low, medium, high, xhigh" };
+  }
+  return { ok: true, value: value as ThinkingLevel };
+}
+
+function parsePlanModelEntry(value: unknown): IParseResult<IPlanModelEntry> {
+  if (typeof value === "string") {
+    const parsedSelector = parsePlanModelSelector(value);
+    if (!parsedSelector.ok) {
+      return parsedSelector;
+    }
+    return {
+      ok: true,
+      value: {
+        model: parsedSelector.value,
+      },
+    };
+  }
+
+  if (!isRecord(value)) {
+    return { ok: false, error: "expected model selector string or { model, thinkingLevel } object" };
+  }
+
+  const knownKeys = new Set(["model", "thinkingLevel"]);
+  for (const key of Object.keys(value)) {
+    if (!knownKeys.has(key)) {
+      return { ok: false, error: `unknown key: ${key}` };
+    }
+  }
+
+  const parsedSelector = parsePlanModelSelector(value.model);
+  if (!parsedSelector.ok) {
+    return { ok: false, error: `model entry: ${parsedSelector.error}` };
+  }
+
+  const entry: IPlanModelEntry = {
+    model: parsedSelector.value,
+  };
+
+  if (value.thinkingLevel !== undefined) {
+    const parsedThinking = parseThinkingLevel(value.thinkingLevel);
+    if (!parsedThinking.ok) {
+      return { ok: false, error: `model entry thinkingLevel: ${parsedThinking.error}` };
+    }
+    entry.thinkingLevel = parsedThinking.value;
+  }
+
+  return { ok: true, value: entry };
+}
+
+function parsePlanModelEntries(value: unknown): IParseResult<IPlanModelEntry[]> {
+  if (!Array.isArray(value)) {
+    return { ok: false, error: "expected array of model entries" };
+  }
+  if (value.length === 0) {
+    return { ok: false, error: "models array must not be empty" };
+  }
+
+  const entries: IPlanModelEntry[] = [];
+  for (const item of value) {
+    const parsed = parsePlanModelEntry(item);
+    if (!parsed.ok) {
+      return parsed;
+    }
+    entries.push(parsed.value);
+  }
+
+  return { ok: true, value: entries };
+}
+
+function parsePlanProfileOverride(name: string, value: unknown): IParseResult<IPlanProfileOverride> {
+  if (!isRecord(value)) {
+    return { ok: false, error: `profiles.${name}: expected object` };
+  }
+
+  const knownKeys = new Set(["match", "models"]);
+  for (const key of Object.keys(value)) {
+    if (!knownKeys.has(key)) {
+      return { ok: false, error: `profiles.${name}: unknown key: ${key}` };
+    }
+  }
+
+  const parsedMatch = parsePlanModelSelector(value.match);
+  if (!parsedMatch.ok) {
+    return { ok: false, error: `profiles.${name}.match: ${parsedMatch.error}` };
+  }
+
+  const profile: IPlanProfileOverride = {
+    match: parsedMatch.value,
+  };
+
+  if (value.models !== undefined) {
+    const parsedModels = parsePlanModelEntries(value.models);
+    if (!parsedModels.ok) {
+      return { ok: false, error: `profiles.${name}.models: ${parsedModels.error}` };
+    }
+    profile.models = parsedModels.value;
+  }
+
+  return { ok: true, value: profile };
+}
+
+function parsePlanProfiles(value: unknown): IParseResult<Record<string, IPlanProfileOverride>> {
+  if (!isRecord(value)) {
+    return { ok: false, error: "profiles: expected object" };
+  }
+
+  const profiles: Record<string, IPlanProfileOverride> = {};
+  for (const [name, entry] of Object.entries(value)) {
+    const parsed = parsePlanProfileOverride(name, entry);
+    if (!parsed.ok) {
+      return parsed;
+    }
+    profiles[name] = parsed.value;
+  }
+
+  return { ok: true, value: profiles };
+}
+
+export function parsePlanPolicy(value: unknown): IParseResult<IPlanPolicy> {
+  if (!isRecord(value)) {
+    return { ok: false, error: "expected object" };
+  }
+
+  const knownKeys = new Set(["models", "profiles"]);
+  for (const key of Object.keys(value)) {
+    if (!knownKeys.has(key)) {
+      return { ok: false, error: `unknown key: ${key}` };
+    }
+  }
+
+  const policy: IPlanPolicy = {};
+
+  if (value.models !== undefined) {
+    const parsedModels = parsePlanModelEntries(value.models);
+    if (!parsedModels.ok) {
+      return { ok: false, error: `models: ${parsedModels.error}` };
+    }
+    policy.models = parsedModels.value;
+  }
+
+  if (value.profiles !== undefined) {
+    const parsedProfiles = parsePlanProfiles(value.profiles);
+    if (!parsedProfiles.ok) {
+      return parsedProfiles;
+    }
+    policy.profiles = parsedProfiles.value;
+  }
+
+  return { ok: true, value: policy };
+}
+
+function readPlanPolicyConfigFile(configPath: string): IReadPlanPolicyResult {
+  try {
+    const raw = readFileSync(configPath, "utf8");
+    const parsed = parsePlanPolicy(JSON.parse(raw) as unknown);
+    if (!parsed.ok) {
+      return { policy: {}, path: configPath, error: `Invalid ${configPath}: ${parsed.error}` };
+    }
+    return { policy: parsed.value, path: configPath };
+  } catch (error) {
+    const message = error instanceof Error ? error.message : String(error);
+    return { policy: {}, path: configPath, error: `Invalid ${configPath}: ${message}` };
+  }
+}
+
+export function readPlanPolicy(cwd: string): IReadPlanPolicyResult {
+  const projectPath = path.join(cwd, PLAN_CONFIG_FILE);
+  if (existsSync(projectPath)) {
+    return readPlanPolicyConfigFile(projectPath);
+  }
+
+  const globalPath = resolveGlobalPlanPolicyPath();
+  if (existsSync(globalPath)) {
+    return readPlanPolicyConfigFile(globalPath);
+  }
+
+  return { policy: {} };
+}
+
+export function findMatchingPlanProfile(
+  profiles: Record<string, IPlanProfileOverride> | undefined,
+  sourceModel: string | undefined,
+): { name: string; override: IPlanProfileOverride } | undefined {
+  if (!profiles || !sourceModel) {
+    return undefined;
+  }
+
+  for (const name of Object.keys(profiles).sort()) {
+    const profile = profiles[name];
+    if (profile && profile.match === sourceModel) {
+      return { name, override: profile };
+    }
+  }
+
+  return undefined;
+}
+
+export function resolveEffectivePlanPolicy(params: {
+	basePolicy: IPlanPolicy;
+	sourceModel: string | undefined;
+}): {
+	profileName: string | undefined;
+	modelEntries: IPlanModelEntry[];
+} {
+  const profile = findMatchingPlanProfile(params.basePolicy.profiles, params.sourceModel);
+  return {
+    profileName: profile?.name,
+    modelEntries: profile?.override.models ?? params.basePolicy.models ?? [],
+  };
+}
+
+function formatModelSelector(model: { provider: string; id: string }): string {
+  return `${model.provider}/${model.id}`;
+}
+
+function splitModelSelector(selector: string): IParseResult<{ provider: string; modelId: string }> {
+  const parsed = parsePlanModelSelector(selector);
+  if (!parsed.ok) {
+    return parsed;
+  }
+
+  const slashIndex = parsed.value.indexOf("/");
+  return {
+    ok: true,
+    value: {
+      provider: parsed.value.slice(0, slashIndex),
+      modelId: parsed.value.slice(slashIndex + 1),
+    },
+  };
+}
+
+function getCurrentModelSelector(ctx: ExtensionContext): string | undefined {
+  return ctx.model ? formatModelSelector(ctx.model) : undefined;
+}
+
+async function tryResolveConfiguredModel(
+  ctx: ExtensionContext,
+  selector: string,
+): Promise<Model<Api> | undefined> {
+  const parsed = splitModelSelector(selector);
+  if (!parsed.ok) {
+    return undefined;
+  }
+
+  const model = ctx.modelRegistry.find(parsed.value.provider, parsed.value.modelId);
+  if (!model) {
+    return undefined;
+  }
+
+  try {
+    const apiKey = await ctx.modelRegistry.getApiKey(model);
+    if (!apiKey) {
+      return undefined;
+    }
+  } catch {
+    return undefined;
+  }
+
+  return model as Model<Api>;
+}
+
+async function resolveConfiguredPlanModel(
+  ctx: ExtensionContext,
+  modelEntries: IPlanModelEntry[],
+): Promise<{ entry: IPlanModelEntry; model: Model<Api> } | undefined> {
+  for (const entry of modelEntries) {
+    const model = await tryResolveConfiguredModel(ctx, entry.model);
+    if (model) {
+      return { entry, model };
+    }
+  }
+
+  return undefined;
+}
+
+function buildPlanPromptOverlay(state: IPlanSessionState): string {
+  const task = state.task?.trim() || "(unknown task)";
+  const model = state.planModel ? `Planning model: ${state.planModel}` : "Planning model: current session model";
+  const artifact = state.artifactPath ? `Artifact path: ${state.artifactPath}` : "Artifact path: unavailable";
+  const modeLabel = state.mode === "architect" ? "Docyrus `/architect` planning mode." : "Docyrus `/plan` mode.";
+  return [
+    modeLabel,
+    PLAN_MODE_SYSTEM_PROMPT,
+    "",
+    buildAskUserProtocolInstructions(),
+    "",
+    "Current planning session:",
+    `Task: ${task}`,
+    model,
+    artifact,
+  ].join("\n");
+}
+
+function buildInitialPlanPrompt(task: string | undefined): string {
+  if (!task) {
+    return [
+      "The user started a planning workflow but did not provide a concrete task yet.",
+      "",
+      "Before drafting any plan, gather the minimum clarification you need from the user.",
+      "If clarification is needed, use the ask_user protocol.",
+    ].join("\n");
+  }
+
+  return [
+    "Create a decision-complete implementation plan for the following task.",
+    "",
+    task,
+  ].join("\n");
+}
+
+function buildInitialArtifactContent(task: string): string {
+  return [
+    "# Planning Session",
+    "",
+    "## Task",
+    task,
+    "",
+    "_Waiting for the first planning response..._",
+    "",
+  ].join("\n");
+}
+
+function getPlanState(ctx: ExtensionContext): IPlanSessionState | undefined {
+  let state: IPlanSessionState | undefined;
+  for (const entry of ctx.sessionManager.getBranch()) {
+    if (entry.type === "custom" && entry.customType === PLAN_STATE_TYPE) {
+      state = entry.data as IPlanSessionState | undefined;
+    }
+  }
+
+  if (state?.active) {
+    return state;
+  }
+
+  return undefined;
+}
+
+function setPlanWidget(ctx: ExtensionContext, state: IPlanSessionState | undefined): void {
+  if (!ctx.hasUI) {
+    return;
+  }
+
+  if (!state?.active) {
+    ctx.ui.setWidget(PLAN_WIDGET_KEY, undefined);
+    return;
+  }
+
+  ctx.ui.setWidget(PLAN_WIDGET_KEY, (_tui, theme) => {
+    const lines = [
+      theme.fg("accent", theme.bold("Plan mode active")),
+      theme.fg("muted", `Model: ${state.planModel ?? "current session model"}`),
+      theme.fg("muted", `Artifact: ${state.artifactPath ?? "(none)"}`),
+      theme.fg("warning", "Exit with /end-plan"),
+    ];
+
+    const container = new Container();
+    container.addChild(new DynamicBorder((value: string) => theme.fg("accent", value)));
+    container.addChild(new Text(lines.join("\n"), 0, 0));
+    return container;
+  });
+}
+
+async function ensureArtifactFile(artifactPath: string, task: string): Promise<void> {
+  await fs.mkdir(path.dirname(artifactPath), { recursive: true });
+  await fs.writeFile(artifactPath, buildInitialArtifactContent(task), "utf8");
+}
+
+function appendPlanState(pi: ExtensionAPI, state: IPlanSessionState): void {
+  pi.appendEntry(PLAN_STATE_TYPE, state);
+}
+
+function clearPlanState(pi: ExtensionAPI): void {
+  pi.appendEntry(PLAN_STATE_TYPE, { active: false });
+}
+
+function getLatestUserEntryId(ctx: ExtensionContext): string | undefined {
+  const branch = ctx.sessionManager.getBranch();
+  for (let index = branch.length - 1; index >= 0; index -= 1) {
+    const entry = branch[index];
+    if (entry.type === "message" && entry.message.role === "user") {
+      return entry.id;
+    }
+  }
+
+  return undefined;
+}
+
+function extractLastAssistantText(messages: Array<{ role?: string; content?: unknown }>): string | null {
+  for (let index = messages.length - 1; index >= 0; index -= 1) {
+    const message = messages[index];
+    if (message?.role !== "assistant") {
+      continue;
+    }
+
+    if (typeof message.content === "string") {
+      return message.content.trim() || null;
+    }
+
+    if (Array.isArray(message.content)) {
+      const text = message.content
+        .filter((block): block is { type: "text"; text: string } => {
+          return Boolean(
+            block &&
+						typeof block === "object" &&
+						"type" in block &&
+						"text" in block &&
+						block.type === "text" &&
+						typeof block.text === "string",
+          );
+        })
+        .map((block) => block.text)
+        .join("\n")
+        .trim();
+
+      return text || null;
+    }
+
+    return null;
+  }
+
+  return null;
+}
+
+async function writePlanArtifactFromEvent(event: AgentEndEvent, ctx: ExtensionContext): Promise<void> {
+  const state = getPlanState(ctx);
+  if (!state?.active || !state.artifactPath || state.mode === "architect") {
+    return;
+  }
+
+  const text = extractLastAssistantText(event.messages ?? []);
+  if (!text || parseAskUserRequestFromText(text)) {
+    return;
+  }
+
+  try {
+    await fs.mkdir(path.dirname(state.artifactPath), { recursive: true });
+    await fs.writeFile(state.artifactPath, `${text.trim()}\n`, "utf8");
+  } catch (error) {
+    if (ctx.hasUI) {
+      const message = error instanceof Error ? error.message : String(error);
+      ctx.ui.notify(`Failed to write plan artifact: ${message}`, "error");
+    }
+  }
+}
+
+export function shouldBlockTodoAction(action: string | undefined): boolean {
+  return !!action && !ALLOWED_TODO_ACTIONS.has(action);
+}
+
+export function isPotentiallyMutatingShellCommand(command: string): boolean {
+  const trimmed = command.trim();
+  if (!trimmed) {
+    return false;
+  }
+
+  const patterns = [
+    /(^|[;&|]\s*)git\s+(add|commit|push|pull|merge|rebase|cherry-pick|reset|clean|restore|checkout|switch|stash|am|apply)\b/i,
+    /(^|[;&|]\s*)(rm|mv|cp|mkdir|touch|install|chmod|chown)\b/i,
+    /(^|[;&|]\s*)(npm|pnpm|yarn|bun)\s+(add|install|remove|unlink|link|update|upgrade|publish|version)\b/i,
+    /(^|[;&|]\s*)(pip|pip3)\s+(install|uninstall)\b/i,
+    /(^|[;&|]\s*)cargo\s+(add|install|remove)\b/i,
+    /(^|[;&|]\s*)go\s+(get|mod\s+tidy)\b/i,
+    /(^|[;&|]\s*)brew\s+(install|upgrade|tap|untap)\b/i,
+    /(^|[;&|]\s*)sed\s+-i\b/i,
+    /(^|[;&|]\s*)perl\s+-i\b/i,
+    /(^|[;&|]\s*)tee\b/i,
+    /(^|[;&|]\s*)docyrus\s+(apps|ds|studio)\s+(create|update|delete|restore|permanent-delete|bulk-create-data-sources|create-field|update-field|delete-field|create-fields-batch|update-fields-batch|delete-fields-batch|create-enums|update-enums|delete-enums)\b/i,
+  ];
+
+  if (patterns.some((pattern) => pattern.test(trimmed))) {
+    return true;
+  }
+
+  if (/\s(?:\d*>>?)(?:\s|$)/.test(trimmed)) {
+    return true;
+  }
+
+  return false;
+}
+
+function buildBlockedToolReason(detail: string): string {
+  return `A planning session is active. ${detail} Exit with /end-plan or /end-architect before implementing.`;
+}
+
+function buildBlockedUserBashResult(detail: string) {
+  return {
+    result: {
+      output: `${buildBlockedToolReason(detail)}\n`,
+      exitCode: 1,
+      cancelled: false,
+      truncated: false,
+    },
+  };
+}
+
+async function applyModelSelector(
+  pi: ExtensionAPI,
+  ctx: ExtensionContext,
+  selector: string | undefined,
+  thinkingLevel: ThinkingLevel | undefined,
+): Promise<boolean> {
+  if (!selector) {
+    if (thinkingLevel) {
+      pi.setThinkingLevel(thinkingLevel);
+    }
+    return true;
+  }
+
+  const parts = splitModelSelector(selector);
+  if (!parts.ok) {
+    if (ctx.hasUI) {
+      ctx.ui.notify(`Invalid plan model selector ${selector}`, "warning");
+    }
+    return false;
+  }
+
+  const model = ctx.modelRegistry.find(parts.value.provider, parts.value.modelId);
+  if (!model) {
+    if (ctx.hasUI) {
+      ctx.ui.notify(`Plan model ${selector} is not available in the current registry`, "warning");
+    }
+    return false;
+  }
+
+  const changed = getCurrentModelSelector(ctx) !== selector;
+  const ok = changed ? await pi.setModel(model) : true;
+  if (!ok) {
+    if (ctx.hasUI) {
+      ctx.ui.notify(`No API key available for plan model ${selector}`, "warning");
+    }
+    return false;
+  }
+
+  if (thinkingLevel) {
+    pi.setThinkingLevel(thinkingLevel);
+  }
+
+  return true;
+}
+
+async function applyResolvedPlanModel(
+  pi: ExtensionAPI,
+  ctx: ExtensionContext,
+  resolved: { entry: IPlanModelEntry; model: Model<Api> } | undefined,
+): Promise<{ planModel?: string; planThinking?: ThinkingLevel }> {
+  if (!resolved) {
+    return {
+      planModel: getCurrentModelSelector(ctx),
+      planThinking: pi.getThinkingLevel(),
+    };
+  }
+
+  const ok = await pi.setModel(resolved.model);
+  if (!ok) {
+    if (ctx.hasUI) {
+      ctx.ui.notify(`No API key available for plan model ${resolved.entry.model}`, "warning");
+    }
+    return {
+      planModel: getCurrentModelSelector(ctx),
+      planThinking: pi.getThinkingLevel(),
+    };
+  }
+
+  if (resolved.entry.thinkingLevel) {
+    pi.setThinkingLevel(resolved.entry.thinkingLevel);
+  }
+
+  return {
+    planModel: formatModelSelector(resolved.model),
+    planThinking: resolved.entry.thinkingLevel ?? pi.getThinkingLevel(),
+  };
+}
+
+async function restoreSourceModel(
+  pi: ExtensionAPI,
+  ctx: ExtensionContext,
+  state: IPlanSessionState | undefined,
+): Promise<void> {
+  if (!state) {
+    return;
+  }
+
+  const restored = await applyModelSelector(pi, ctx, state.sourceModel, state.sourceThinking);
+  if (!restored && state.sourceThinking) {
+    pi.setThinkingLevel(state.sourceThinking);
+  }
+}
+
+async function syncPlanState(pi: ExtensionAPI, ctx: ExtensionContext): Promise<void> {
+  const nextState = getPlanState(ctx);
+  const previousState = currentPlanState;
+  currentPlanState = nextState;
+  setPlanWidget(ctx, nextState);
+
+  if (nextState?.active) {
+    const currentSelector = getCurrentModelSelector(ctx);
+    const currentThinking = pi.getThinkingLevel();
+    if (nextState.planModel !== currentSelector || nextState.planThinking !== currentThinking) {
+      await applyModelSelector(pi, ctx, nextState.planModel, nextState.planThinking);
+    }
+    return;
+  }
+
+  if (previousState?.active) {
+    await restoreSourceModel(pi, ctx, previousState);
+  }
+}
+
+function formatModelChain(entries: IPlanModelEntry[]): string {
+  if (entries.length === 0) {
+    return "(none)";
+  }
+  return entries
+    .map((entry) => entry.thinkingLevel ? `${entry.model} (thinking: ${entry.thinkingLevel})` : entry.model)
+    .join(" -> ");
+}
+
+export async function startPlanningWorkflow(params: IPlanningStartParams): Promise<boolean> {
+  const { pi, ctx, mode, artifactPath, initialPrompt } = params;
+  if (getPlanState(ctx)) {
+    if (ctx.hasUI) {
+      ctx.ui.notify("A planning session is already active. End it before starting another one.", "warning");
+    }
+    return false;
+  }
+
+  const sourceModel = getCurrentModelSelector(ctx);
+  const sourceThinking = pi.getThinkingLevel();
+  const policyResult = readPlanPolicy(ctx.cwd);
+  if (policyResult.error && ctx.hasUI) {
+    ctx.ui.notify(policyResult.error, "warning");
+  }
+
+  const { profileName, modelEntries } = resolveEffectivePlanPolicy({
+    basePolicy: policyResult.policy,
+    sourceModel,
+  });
+  const resolvedPlanModel = await resolveConfiguredPlanModel(ctx, modelEntries);
+  if (!resolvedPlanModel && modelEntries.length > 0) {
+    ctx.ui.notify(
+      `No configured planning models could be resolved (${formatModelChain(modelEntries)}). Staying on the current model.`,
+      "warning",
+    );
+  }
+
+  try {
+    await ensureArtifactFile(artifactPath, params.task ?? mode);
+  } catch (error) {
+    if (ctx.hasUI) {
+      const message = error instanceof Error ? error.message : String(error);
+      ctx.ui.notify(`Failed to create plan artifact ${artifactPath}: ${message}`, "error");
+    }
+    return false;
+  }
+
+  let originId = ctx.sessionManager.getLeafId() ?? undefined;
+  if (!originId) {
+    pi.appendEntry(PLAN_ANCHOR_TYPE, { createdAt: new Date().toISOString() });
+    originId = ctx.sessionManager.getLeafId() ?? undefined;
+  }
+
+  if (!originId) {
+    if (ctx.hasUI) {
+      ctx.ui.notify("Failed to determine the current branch position for the planning workflow.", "error");
+    }
+    return false;
+  }
+
+  let branchStartTargetId = getLatestUserEntryId(ctx);
+  if (branchStartTargetId && branchStartTargetId === originId) {
+    pi.appendEntry(PLAN_ANCHOR_TYPE, { createdAt: new Date().toISOString() });
+    originId = ctx.sessionManager.getLeafId() ?? originId;
+  }
+
+  if (branchStartTargetId && branchStartTargetId !== originId) {
+    try {
+      const result = await ctx.navigateTree(branchStartTargetId, {
+        summarize: false,
+        label: PLAN_BRANCH_LABEL,
+      });
+      if (result.cancelled) {
+        return false;
+      }
+    } catch (error) {
+      if (ctx.hasUI) {
+        const message = error instanceof Error ? error.message : String(error);
+        ctx.ui.notify(`Failed to start the planning branch: ${message}`, "error");
+      }
+      return false;
+    }
+  }
+
+  const branchStartId = ctx.sessionManager.getLeafId() ?? branchStartTargetId;
+  const appliedPlanModel = await applyResolvedPlanModel(pi, ctx, resolvedPlanModel);
+  const state: IPlanSessionState = {
+    active: true,
+    mode,
+    originId,
+    branchStartId,
+    artifactPath,
+    task: params.task,
+    sourceModel,
+    sourceThinking,
+    planModel: appliedPlanModel.planModel,
+    planThinking: appliedPlanModel.planThinking,
+    startedAt: new Date().toISOString(),
+  };
+
+  appendPlanState(pi, state);
+  currentPlanState = state;
+  setPlanWidget(ctx, state);
+
+  if (ctx.hasUI) {
+    const profileHint = profileName ? ` (profile: ${profileName})` : "";
+    const modeLabel = mode === "architect" ? "Architect mode" : "Plan mode";
+    ctx.ui.notify(`${modeLabel} active${profileHint}. Writing updates to ${artifactPath}`, "info");
+  }
+
+  pi.sendUserMessage(initialPrompt);
+  return true;
+}
+
+async function planHandler(pi: ExtensionAPI, ctx: ExtensionCommandContext, args: string): Promise<void> {
+  const task = parsePlanTask(args) ?? undefined;
+  const artifactPath = createPlanArtifactPath({
+    cwd: ctx.cwd,
+    task: task ?? "plan",
+    date: new Date(),
+  });
+
+  await startPlanningWorkflow({
+    pi,
+    ctx,
+    mode: "plan",
+    task,
+    artifactPath,
+    initialPrompt: buildInitialPlanPrompt(task),
+  });
+}
+
+export async function endPlanningWorkflow(pi: ExtensionAPI, ctx: ExtensionCommandContext): Promise<void> {
+  const state = getPlanState(ctx);
+  if (!state?.active || !state.originId) {
+    if (ctx.hasUI) {
+      ctx.ui.notify("No active planning session was found on the current branch.", "info");
+    }
+    return;
+  }
+
+  clearPlanState(pi);
+
+  try {
+    const result = await ctx.navigateTree(state.originId, {
+      summarize: true,
+      customInstructions: END_PLAN_SUMMARY_PROMPT,
+      replaceInstructions: true,
+      label: PLAN_BRANCH_LABEL,
+    });
+
+    if (result.cancelled) {
+      currentPlanState = undefined;
+      setPlanWidget(ctx, undefined);
+      await restoreSourceModel(pi, ctx, state);
+      if (ctx.hasUI) {
+        ctx.ui.notify("Planning session ended in place after navigation was cancelled.", "info");
+      }
+      return;
+    }
+  } catch (error) {
+    currentPlanState = undefined;
+    setPlanWidget(ctx, undefined);
+    await restoreSourceModel(pi, ctx, state);
+    if (ctx.hasUI) {
+      const message = error instanceof Error ? error.message : String(error);
+      ctx.ui.notify(`Failed to return from the planning branch: ${message}`, "error");
+    }
+    return;
+  }
+
+  currentPlanState = undefined;
+  setPlanWidget(ctx, undefined);
+  await restoreSourceModel(pi, ctx, state);
+
+  if (ctx.hasUI) {
+    const artifactSuffix = state.artifactPath ? ` Artifact: ${state.artifactPath}` : "";
+    ctx.ui.notify(`Planning session ended.${artifactSuffix}`, "info");
+  }
+}
+
+async function collectAskUserResponse(ctx: ExtensionContext, request: { title: string; message: string; questions: IAskUserQuestion[] }): Promise<IAskUserResponse | undefined> {
+  if (!ctx.hasUI) {
+    return undefined;
+  }
+
+  const answers: Record<string, string | boolean | string[]> = {};
+
+  for (const question of request.questions) {
+    if (question.type === "boolean") {
+      const confirmed = await ctx.ui.confirm(question.label, question.description ?? request.message);
+      answers[question.id] = confirmed;
+      continue;
+    }
+
+    if (question.type === "singleSelect" && question.options && question.options.length > 0) {
+      const selectedLabel = await ctx.ui.select(
+        question.label,
+        question.options.map((option) => option.label),
+      );
+      const selected = question.options.find((option) => option.label === selectedLabel);
+      if (selected) {
+        answers[question.id] = selected.value;
+      }
+      continue;
+    }
+
+    const prefill = Array.isArray(question.defaultValue)
+      ? question.defaultValue.join(", ")
+      : typeof question.defaultValue === "string"
+        ? question.defaultValue
+        : "";
+    const input = await ctx.ui.editor(question.label, prefill);
+    if (input === undefined) {
+      return {
+        action: "cancel",
+        answers,
+      };
+    }
+
+    if (question.type === "multiSelect") {
+      answers[question.id] = input
+        .split(",")
+        .map((value) => value.trim())
+        .filter(Boolean);
+    } else {
+      answers[question.id] = input.trim();
+    }
+  }
+
+  return {
+    action: "submit",
+    answers,
+  };
+}
+
+async function planPolicyHandler(_pi: ExtensionAPI, ctx: ExtensionCommandContext): Promise<void> {
+  const branchState = getPlanState(ctx);
+  const sourceModel = branchState?.sourceModel ?? getCurrentModelSelector(ctx);
+  const currentModel = getCurrentModelSelector(ctx);
+  const policyResult = readPlanPolicy(ctx.cwd);
+  if (policyResult.error) {
+    ctx.ui.notify(policyResult.error, "warning");
+  }
+
+  const { profileName, modelEntries } = resolveEffectivePlanPolicy({
+    basePolicy: policyResult.policy,
+    sourceModel,
+  });
+  const resolved = await resolveConfiguredPlanModel(ctx, modelEntries);
+  const resolvedSelector = resolved ? formatModelSelector(resolved.model) : branchState?.planModel ?? currentModel ?? "(unresolved)";
+  const resolvedThinking = resolved?.entry.thinkingLevel ?? branchState?.planThinking;
+
+  const lines = [
+    `status: ${branchState?.active ? "active" : "inactive"}`,
+    `config: ${policyResult.path ?? "(none)"}`,
+    `source model: ${sourceModel ?? "unknown"}`,
+    `current model: ${currentModel ?? "unknown"}`,
+    `profile: ${profileName ?? "none"}`,
+    `configured models: ${formatModelChain(modelEntries)}`,
+    `resolved planning model: ${resolvedThinking ? `${resolvedSelector} (thinking: ${resolvedThinking})` : resolvedSelector}`,
+  ];
+
+  if (branchState?.artifactPath) {
+    lines.push(`artifact: ${branchState.artifactPath}`);
+  }
+
+  ctx.ui.notify(lines.join("\n"), "info");
+}
+
+function handleToolCallDuringPlan(event: ToolCallEvent, ctx: ExtensionContext): { block: boolean; reason: string } | undefined {
+  if (!getPlanState(ctx)?.active) {
+    return undefined;
+  }
+
+  if (event.toolName === "edit" || event.toolName === "write") {
+    return {
+      block: true,
+      reason: buildBlockedToolReason(`The ${event.toolName} tool is disabled during planning.`),
+    };
+  }
+
+  if (event.toolName === "bash" && isRecord(event.input) && typeof event.input.command === "string") {
+    if (isPotentiallyMutatingShellCommand(event.input.command)) {
+      return {
+        block: true,
+        reason: buildBlockedToolReason(`Mutating bash commands are disabled during planning (${event.input.command}).`),
+      };
+    }
+  }
+
+  if (event.toolName === "todo" && isRecord(event.input) && typeof event.input.action === "string") {
+    if (shouldBlockTodoAction(event.input.action)) {
+      return {
+        block: true,
+        reason: buildBlockedToolReason(`Todo action "${event.input.action}" is disabled during planning.`),
+      };
+    }
+  }
+
+  return undefined;
+}
+
+function handleUserBashDuringPlan(event: UserBashEvent, ctx: ExtensionContext) {
+  if (!getPlanState(ctx)?.active) {
+    return undefined;
+  }
+
+  if (!isPotentiallyMutatingShellCommand(event.command)) {
+    return undefined;
+  }
+
+  return buildBlockedUserBashResult(`Mutating shell commands are disabled during planning (${event.command}).`);
+}
+
+export default function planExtension(pi: ExtensionAPI) {
+  pi.registerCommand("plan", {
+    description: "Start a planning-only branch with an optional dedicated planning model",
+    handler: async(args, ctx) => {
+      await planHandler(pi, ctx, args);
+    },
+  });
+
+  pi.registerCommand("end-plan", {
+    description: "Leave plan mode, summarize the plan branch, and return to the original branch",
+    handler: async(_args, ctx) => {
+      await endPlanningWorkflow(pi, ctx);
+    },
+  });
+
+  pi.registerCommand("end-architect", {
+    description: "Leave architect mode, summarize the architect branch, and return to the original branch",
+    handler: async(_args, ctx) => {
+      await endPlanningWorkflow(pi, ctx);
+    },
+  });
+
+  pi.registerCommand("plan-policy", {
+    description: "Show effective planning-model policy and the resolved planning model",
+    handler: async(_args, ctx) => {
+      await planPolicyHandler(pi, ctx);
+    },
+  });
+
+  pi.on("before_agent_start", (event, ctx) => {
+    const state = getPlanState(ctx);
+    if (!state?.active) {
+      return;
+    }
+
+    return {
+      systemPrompt: [event.systemPrompt, buildPlanPromptOverlay(state)].filter(Boolean).join("\n\n"),
+    };
+  });
+
+  pi.on("tool_call", (event, ctx) => {
+    return handleToolCallDuringPlan(event, ctx);
+  });
+
+  pi.on("user_bash", (event, ctx) => {
+    return handleUserBashDuringPlan(event, ctx);
+  });
+
+  pi.on("agent_end", async(event, ctx) => {
+    const state = getPlanState(ctx);
+    if (state?.active) {
+      const text = extractLastAssistantText(event.messages ?? []);
+      const askUserRequest = text ? parseAskUserRequestFromText(text) : undefined;
+      if (askUserRequest && ctx.hasUI) {
+        const response = await collectAskUserResponse(ctx, askUserRequest);
+        if (response) {
+          pi.sendUserMessage(formatAskUserResponsePrompt(response));
+        }
+        return;
+      }
+    }
+
+    await writePlanArtifactFromEvent(event, ctx);
+  });
+
+  pi.on("session_start", async(_event, ctx) => {
+    await syncPlanState(pi, ctx);
+  });
+
+  pi.on("session_switch", async(_event, ctx) => {
+    await syncPlanState(pi, ctx);
+  });
+
+  pi.on("session_tree", async(_event, ctx) => {
+    await syncPlanState(pi, ctx);
+  });
+
+  pi.on("session_shutdown", async(_event, ctx) => {
+    currentPlanState = undefined;
+    setPlanWidget(ctx, undefined);
+  });
+}