@tduyng/pi-prewalk 0.0.0-stage → 0.0.1

package/src/index.ts

@@ -0,0 +1,925 @@
+import type {
+    ExtensionAPI,
+    ExtensionContext,
+    ModelRoute,
+    ModelRouteRequest,
+    ToolDefinition
+} from '@earendil-works/pi-coding-agent'
+import {
+    contextForModel,
+    continuationNudge,
+    ESCALATION_BAR,
+    makeSpecPredicate,
+    modelRef,
+    parseArgs,
+    planProgress,
+    PLAN_DEMAND_REASON,
+    PLANNING_BRIEF,
+    PLAN_TOOL_NAME,
+    READ_ONLY_TOOLS,
+    resolveModelSpec,
+    shortName,
+    STATE_TYPE,
+    statusBadge,
+    statusLabel,
+    stepForCommand,
+    stepForPath,
+    swapDecision,
+    MAX_PLAN_FACTS,
+    MAX_PLAN_STEPS,
+    MAX_WALL_CHARS,
+    NUDGE_PREFIX,
+    SKILL_BLOCK_PREFIX,
+    shouldAsk,
+    isScratchPath,
+    PLANNER_GRACE_CALLS,
+    renderHandoff,
+    reviewPrompt,
+    wallBrief,
+    type ForceRoute,
+    type PrewalkState
+} from './logic.js'
+import { bashWritesFiles } from './bash-writes.js'
+import { readEffectivePrewalkConfig, writePrewalkConfig, type PrewalkConfig, type PrewalkMode } from './config.js'
+
+const MODE_LABEL: Record<PrewalkMode, string> = {
+    plan: 'every task plans',
+    off: 'the cheap model, no plan'
+}
+
+/** One decision: the model the task starts on. */
+type Routed = { route: ForceRoute }
+
+const ROUTER_PROVIDER = 'prewalk'
+const ROUTER_ID = 'auto'
+const ROUTER_REF = `${ROUTER_PROVIDER}/${ROUTER_ID}`
+
+import { type Infer, Validator } from 'ata-validator'
+import { t } from 'ata-validator/t'
+
+type SessionModel = NonNullable<ExtensionContext['model']>
+
+/**
+ * pi's api id for a virtual catalog entry. pi keeps that constant and its guard private, so a chain
+ * naming `prewalk/auto` would otherwise resolve and route() would return a virtual model as its
+ * target, which pi rejects with "which is not a physical model".
+ */
+const VIRTUAL_API = 'pi-virtual'
+
+/** The catalog prewalk routes to: everything except the virtual entries, which pi cannot dispatch to. */
+function physicalModels(ctx: ExtensionContext): readonly SessionModel[] {
+    return ctx.modelRegistry.getAvailable().filter((model) => model.api !== VIRTUAL_API)
+}
+
+/**
+ * A chain that resolves to nothing is a configuration problem, so prewalk stops routing rather than
+ * ending the turn. The model already in play wins.
+ */
+function heldModel(
+    ctx: ExtensionContext,
+    inPlay: ModelRouteRequest['previous'],
+    reason: string,
+    level: ModelRouteRequest['thinkingLevel']
+): ModelRoute {
+    ctx.ui.notify(`${reason}; this task runs on the model already in play`, 'error')
+    const held = inPlay?.model ?? physicalModels(ctx)[0]
+    if (!held) throw new Error('prewalk: no model is available to hold the task on')
+    return { model: held, thinkingLevel: inPlay?.thinkingLevel ?? level }
+}
+
+const PLAN_STEP_SCHEMA = t.object({
+    file: t.string({ maxLength: 4_000 }),
+    change: t.string({ maxLength: 8_000 }),
+    anchor: t.optional(t.string({ maxLength: 8_000 })),
+    verify: t.optional(t.string({ maxLength: 2_000 }))
+})
+
+/** The shape a `prewalk` session entry holds, so the read that restores it is bounded. */
+const STORED_ENTRY_SCHEMA = t.object({
+    executor: t.optional(t.string()),
+    planner: t.optional(t.string()),
+    swapped: t.optional(t.boolean()),
+    inFlight: t.optional(t.boolean()),
+    reviewing: t.optional(t.boolean()),
+    planDemanded: t.optional(t.boolean()),
+    forced: t.optional(t.union([t.literal('plan'), t.literal('run')])),
+    stepIndex: t.optional(t.integer({ minimum: 0 })),
+    failures: t.optional(t.integer({ minimum: 0 })),
+    escalations: t.optional(t.integer({ minimum: 0 })),
+    wall: t.optional(t.union([t.string({ maxLength: MAX_WALL_CHARS }), t.null()])),
+    plan: t.optional(
+        t.object({
+            summary: t.string({ maxLength: 8_000 }),
+            facts: t.array(t.string(), { maxItems: MAX_PLAN_FACTS }),
+            steps: t.array(PLAN_STEP_SCHEMA, { maxItems: MAX_PLAN_STEPS }),
+            check: t.string({ maxLength: 2_000 }),
+            outOfScope: t.array(t.string(), { maxItems: MAX_PLAN_FACTS })
+        })
+    )
+})
+const isStoredEntrySchema = new Validator(STORED_ENTRY_SCHEMA).isValidObject
+
+const TOOL_INPUT_SCHEMA = t.object({
+    path: t.optional(t.string()),
+    file: t.optional(t.string()),
+    command: t.optional(t.string())
+})
+const isToolInputSchema = new Validator(TOOL_INPUT_SCHEMA).isValidObject
+
+const PLAN_PARAMETERS = t.object({
+    summary: t.string({ description: 'One-line description of what the plan accomplishes.' }),
+    facts: t.array(t.string(), {
+        maxItems: MAX_PLAN_FACTS,
+        description:
+            'Established during exploration: findings the executor cannot re-derive (file:line, constraints, commands). Each string is one fact.'
+    }),
+    steps: t.array(
+        t.object({
+            file: t.string({ description: 'Repo-relative file path, e.g. src/index.ts.' }),
+            change: t.string({ description: 'What to change in this file.' }),
+            anchor: t.optional(
+                t.string({
+                    description:
+                        'Exact text to match for the edit. Prevents the executor from rewriting the wrong site.'
+                })
+            ),
+            verify: t.optional(
+                t.string({
+                    description: 'Command to run after this step, e.g. pnpm vitest run src/foo.test.ts.'
+                })
+            )
+        }),
+        { maxItems: MAX_PLAN_STEPS, description: `Steps in order. No more than ${MAX_PLAN_STEPS}.` }
+    ),
+    check: t.string({ description: 'Command to run after every step is done, e.g. pnpm check.' }),
+    outOfScope: t.array(t.string(), {
+        maxItems: MAX_PLAN_FACTS,
+        description: 'Files or areas the executor must not touch.'
+    })
+})
+const isPlanParameters = new Validator(PLAN_PARAMETERS).isValidObject
+type PlanParams = Infer<typeof PLAN_PARAMETERS>
+
+type PlanToolDefinition = ToolDefinition<typeof PLAN_PARAMETERS> & { defaultActive?: boolean }
+
+export default function prewalkExtension(pi: ExtensionAPI): void {
+    let state: PrewalkState | null = null
+    let config: PrewalkConfig = readEffectivePrewalkConfig()
+    let stallNudged = false
+    let completionNudged = false
+    let reviewSpent = false
+    let postPlanCalls: number | null = null
+    let demandedThisTurn = false
+    let models: readonly SessionModel[] = []
+
+    /**
+     * The prompt that opened a task, set only by a prompt typed while the session was idle, or by
+     * `/prewalk` with a task. `route()` consumes it, so one prompt routes once. pi's
+     * `reason: "user"` cannot carry this: it also fires for a steer and a follow-up, which are
+     * turns inside the task already in flight.
+     */
+    let prompt: string | null = null
+    let steer: string | null = null
+
+    /**
+     * A prompt too short to be a task is not a task boundary. It holds whatever is in play rather
+     * than opening a task on the cheap tier with nothing planned to execute.
+     */
+    function recordPrompt(text: string): void {
+        if (shouldAsk(text)) prompt = text
+    }
+
+    function persist(): void {
+        pi.appendEntry(STATE_TYPE, state)
+    }
+
+    function updateStatus(ctx: ExtensionContext): void {
+        if (!ctx.hasUI) return
+        if (!state) {
+            ctx.ui.setStatus('prewalk', undefined)
+            return
+        }
+        ctx.ui.setStatus('prewalk', ctx.ui.theme.fg('accent', statusBadge(state)))
+    }
+
+    function chainModel(ctx: ExtensionContext, chain: readonly string[]): SessionModel | null {
+        for (const spec of chain) {
+            const result = resolveModelSpec(physicalModels(ctx), spec)
+            if ('error' in result) continue
+            return result.model
+        }
+        return null
+    }
+
+    function executorModel(ctx: ExtensionContext): SessionModel | null {
+        return chainModel(ctx, config.executor)
+    }
+
+    async function reconstructState(ctx: ExtensionContext): Promise<void> {
+        state = null
+        for (const entry of ctx.sessionManager.getBranch()) {
+            if (entry.type !== 'custom' || entry.customType !== STATE_TYPE) continue
+            const data: unknown = entry.data
+            if (!isStoredEntrySchema(data)) {
+                state = null
+                continue
+            }
+            const plan = data.plan ?? null
+            state = {
+                executor: data.executor ?? '',
+                planner: data.planner ?? null,
+                swapped: data.swapped === true && plan !== null,
+                inFlight: data.inFlight === true,
+                plan,
+                planDemanded: data.planDemanded === true,
+                forced: data.forced ?? null,
+                stepIndex: data.stepIndex ?? 0,
+                failures: data.failures ?? 0,
+                escalations: data.escalations ?? 0,
+                wall: data.wall ?? null,
+                reviewing: data.reviewing === true
+            }
+        }
+        models = physicalModels(ctx)
+        config = readEffectivePrewalkConfig()
+        prompt = null
+        steer = null
+        stallNudged = false
+        completionNudged = false
+        reviewSpent = false
+        postPlanCalls = null
+        updateStatus(ctx)
+        await claimRouter(ctx, ctx.model)
+    }
+
+    function findModel(ctx: ExtensionContext, spec: string): SessionModel | null {
+        const result = resolveModelSpec(physicalModels(ctx), spec)
+        return 'error' in result ? null : result.model
+    }
+
+    function resolveModel(ctx: ExtensionContext, spec: string): SessionModel | null {
+        const result = resolveModelSpec(physicalModels(ctx), spec)
+        if ('error' in result) {
+            ctx.ui.notify(`prewalk: ${result.error}`, 'error')
+            return null
+        }
+        return result.model
+    }
+
+    /**
+     * Every task routes through the router, so a session model that is not the router is a
+     * session where nothing routes, for the rest of its life, and it happens without a word:
+     * one press of the model cycle key leaves prewalk silent. Armed means prewalk owns the
+     * selection, so it takes it back and says so. `/prewalk off` is how a physical model is
+     * kept, and it is the mode that asks nothing and routes nothing.
+     */
+    async function claimRouter(ctx: ExtensionContext, selected: SessionModel | undefined): Promise<void> {
+        if (state === null || config.mode === 'off' || selected === undefined) return
+        if (selected.provider === ROUTER_PROVIDER && selected.id === ROUTER_ID) return
+        if (!(await selectRouter(ctx))) return
+        ctx.ui.notify(
+            `prewalk: ${modelRef(selected)} is not ${ROUTER_REF}, so no task would route. Back on it; /prewalk off keeps the model.`,
+            'warning'
+        )
+    }
+
+    /**
+     * The session's model is the router, so nothing has to be remembered about the swap. Arming only
+     * has to name the two chains and select the router.
+     */
+    async function selectRouter(ctx: ExtensionContext): Promise<boolean> {
+        const router = ctx.modelRegistry.find(ROUTER_PROVIDER, ROUTER_ID)
+        if (!router) {
+            ctx.ui.notify(`prewalk: ${ROUTER_REF} is not registered`, 'error')
+            return false
+        }
+        return pi.setModel(router)
+    }
+
+    function ensureState(ctx: ExtensionContext): PrewalkState | null {
+        if (state) return state
+        const executor = executorModel(ctx)
+        if (!executor) return null
+        const planner = chainModel(ctx, config.planner)
+        state = {
+            executor: modelRef(executor),
+            planner: planner ? modelRef(planner) : null,
+            swapped: false,
+            inFlight: false,
+            forced: null,
+            plan: null,
+            planDemanded: false,
+            stepIndex: 0,
+            failures: 0,
+            escalations: 0,
+            wall: null,
+            reviewing: false
+        }
+        persist()
+        updateStatus(ctx)
+        setPlanTool()
+        return state
+    }
+
+    function disarm(ctx: ExtensionContext): void {
+        state = null
+        prompt = null
+        persist()
+        updateStatus(ctx)
+    }
+
+    function startTask(ctx: ExtensionContext, route_: ForceRoute): void {
+        if (!state) return
+        // A steer typed at the end of a task belongs to that task, not to the next one.
+        steer = null
+        reviewSpent = false
+        postPlanCalls = null
+        state = {
+            ...state,
+            swapped: route_ !== 'plan',
+            inFlight: true,
+            plan: null,
+            planDemanded: route_ !== 'plan',
+            stepIndex: 0,
+            failures: 0,
+            escalations: 0,
+            wall: null,
+            reviewing: false
+        }
+        persist()
+        updateStatus(ctx)
+        setPlanTool()
+    }
+
+    /** The executor takes over on the next request, which `route()` reads off `swapped`. */
+    function markSwapped(ctx: ExtensionContext): void {
+        if (!state || state.swapped) return
+        postPlanCalls = null
+        state = { ...state, swapped: true, failures: 0, wall: null }
+        persist()
+        updateStatus(ctx)
+    }
+
+    function startReview(ctx: ExtensionContext): void {
+        if (!state?.plan) return
+        reviewSpent = true
+        const text = reviewPrompt(state.plan, state.stepIndex)
+        state = { ...state, swapped: false, reviewing: true, inFlight: true, failures: 0, wall: null }
+        persist()
+        updateStatus(ctx)
+        setPlanTool()
+        pi.sendUserMessage(text, { deliverAs: 'steer' })
+    }
+
+    async function arm(
+        ctx: ExtensionContext,
+        args: { planner: string | null; executor: string | null; task: string; forced?: ForceRoute }
+    ): Promise<void> {
+        const executorSpec = args.executor ?? config.executor[0] ?? state?.executor
+        if (!executorSpec) {
+            ctx.ui.notify(
+                'prewalk: no executor named yet.\n\nUsage: /prewalk executor <model> [more fallbacks], or /prewalk <planner> <executor> [task]',
+                'error'
+            )
+            return
+        }
+        const executor = resolveModel(ctx, executorSpec)
+        if (!executor) return
+        const executorRef = modelRef(executor)
+
+        const plannerSpec = args.planner ?? config.planner[0] ?? state?.planner
+        if (!plannerSpec) {
+            ctx.ui.notify('prewalk: no planner named yet.\n\nUsage: /prewalk planner <model> [more fallbacks]', 'error')
+            return
+        }
+        const planner = resolveModel(ctx, plannerSpec)
+        if (!planner) return
+        const plannerRef = modelRef(planner)
+        if (plannerRef === executorRef) {
+            ctx.ui.notify(
+                `prewalk: planner and executor are both ${shortName(executorRef)}; name a different planner`,
+                'error'
+            )
+            return
+        }
+        state = {
+            executor: executorRef,
+            planner: plannerRef,
+            swapped: false,
+            inFlight: false,
+            forced: args.forced ?? null,
+            plan: null,
+            planDemanded: false,
+            stepIndex: 0,
+            failures: 0,
+            escalations: 0,
+            wall: null,
+            reviewing: false
+        }
+        // A named role replaces its chain, a bare re-arm keeps the fallbacks the file already holds.
+        if (args.planner !== null || args.executor !== null) {
+            config = {
+                ...config,
+                planner: args.planner === null ? config.planner : [plannerRef],
+                executor: args.executor === null ? config.executor : [executorRef]
+            }
+            writePrewalkConfig(config)
+        }
+        persist()
+        updateStatus(ctx)
+        setPlanTool()
+        if (!(await selectRouter(ctx))) return
+
+        if (args.task) {
+            recordPrompt(args.task)
+            pi.sendUserMessage(args.task, ctx.isIdle() ? undefined : { deliverAs: 'followUp' })
+            ctx.ui.notify(
+                `prewalk: task started on ${shortName(args.forced === 'run' ? executorRef : plannerRef)}`,
+                'info'
+            )
+            return
+        }
+
+        if (args.forced) {
+            const on = args.forced === 'run' ? shortName(executorRef) : shortName(plannerRef)
+            ctx.ui.notify(`prewalk: the next task ${args.forced === 'run' ? 'runs on' : 'plans on'} ${on}`, 'info')
+            return
+        }
+
+        ctx.ui.notify(
+            `prewalk: ${MODE_LABEL[config.mode]}. Plan on ${shortName(plannerRef)}, run on ${shortName(executorRef)}. /prewalk off stops it`,
+            'info'
+        )
+    }
+
+    function chainReport(ctx: ExtensionContext, role: 'planner' | 'executor', chain: readonly string[]): string {
+        const lines = chain.map((spec) => {
+            const model = findModel(ctx, spec)
+            const mark = model === null ? 'unavailable' : modelRef(model) === ROUTER_REF ? 'router' : 'ready'
+            return `  ${spec}  ${mark}  ${model === null ? '' : priceLine(model)}`
+        })
+        return `${role}: ${chain.length === 0 ? '(none named)' : ''}\n${lines.join('\n')}`
+    }
+
+    /** Both chains with what resolves and what it costs, and the mode, which is what `/prewalk models` is for. */
+    function configReport(ctx: ExtensionContext): string {
+        const chains = `${chainReport(ctx, 'planner', config.planner)}\n${chainReport(ctx, 'executor', config.executor)}`
+        return `${chains}\nmode: ${MODE_LABEL[config.mode]}`
+    }
+
+    function setPlanTool(): void {
+        const enabled = state !== null && config.mode === 'plan'
+        const active = pi.getActiveTools()
+        if (active.length === 0) return
+        if (active.includes(PLAN_TOOL_NAME) === enabled) return
+        pi.setActiveTools(enabled ? [...active, PLAN_TOOL_NAME] : active.filter((name) => name !== PLAN_TOOL_NAME))
+    }
+
+    /**
+     * Nothing is in play: no task, no state, and no model on the request. Work that has not been
+     * planned belongs to the planner, so a session with neither starts there rather than on the
+     * cheap tier.
+     */
+    function coldStart(ctx: ExtensionContext): SessionModel {
+        const planned = chainModel(ctx, config.planner)
+        if (planned) return planned
+        const first = physicalModels(ctx)[0]
+        if (!first) throw new Error('prewalk: no planner in the chain resolves, and no model is available')
+        return first
+    }
+
+    /**
+     * One decision, and the phase answers it: a task plans and the cheap tier runs the steps of
+     * the plan it was given. The score that used to pick between the two read 0.548 AUC against
+     * how long a turn ran, so it did not separate, and it is gone.
+     */
+    function decideRoute(): Routed {
+        const forced = state?.forced ?? null
+        if (forced !== null && state) state = { ...state, forced: null }
+        if (forced !== null) return { route: forced }
+        if (config.mode === 'off') return { route: 'run' }
+        return { route: 'plan' }
+    }
+
+    /**
+     * A recorded wall is the executor saying it cannot finish this alone, which is the one place
+     * where a stuck cheap model hands the task back without asking. One failed tool result is the
+     * evidence, and the task buys two upgrades, because each swap re-reads the whole prefix at the
+     * new model's rate.
+     */
+    async function escalationRoute(
+        ctx: ExtensionContext,
+        inPlay: ModelRouteRequest['previous'],
+        level: ModelRouteRequest['thinkingLevel']
+    ): Promise<ModelRoute | null> {
+        if (config.mode !== 'plan' || state === null) return null
+        if (state.swapped === false || state.wall === null) return null
+        if (state.escalations >= ESCALATION_BAR.upgrades) return null
+        const planned = chainModel(ctx, config.planner)
+        if (!planned)
+            return heldModel(ctx, inPlay, 'prewalk: no planner in the chain resolves, or has an API key', level)
+        state = {
+            ...state,
+            planner: modelRef(planned),
+            swapped: false,
+            escalations: state.escalations + 1,
+            planDemanded: true
+        }
+        persist()
+        updateStatus(ctx)
+        setPlanTool()
+        ctx.ui.notify(
+            `prewalk: the executor is stuck; this task is on ${shortName(modelRef(planned))} until its next edit`,
+            'info'
+        )
+        return { model: planned, thinkingLevel: level }
+    }
+
+    /**
+     * A question typed while the executor is mid-task wants the frontier model, and pi can carry
+     * it under any request reason, so this runs whatever the reason says instead of only reading
+     * continuations: the task goes back to the planner until its next edit hands it over again.
+     * The steer is spent the first time it is read, whether or not the planner is already holding
+     * the question, so it cannot fire on a later swap. No routing call: the phase answers it.
+     */
+    async function steerRoute(
+        ctx: ExtensionContext,
+        executor: SessionModel | null,
+        level: ModelRouteRequest['thinkingLevel']
+    ): Promise<ModelRoute | null> {
+        if (steer === null) return null
+        const text = steer
+        steer = null
+        if (executor === null || state === null) return null
+        if (!shouldAsk(text)) return null
+        const planned = chainModel(ctx, config.planner)
+        if (!planned) return null
+        state = { ...state, planner: modelRef(planned), swapped: false }
+        persist()
+        updateStatus(ctx)
+        setPlanTool()
+        ctx.ui.notify(`prewalk: that question is on ${shortName(modelRef(planned))} until the next edit`, 'info')
+        return { model: planned, thinkingLevel: level }
+    }
+
+    /** What a request costs and the context it has to work in, the two columns a role choice turns on. */
+    function priceLine(model: SessionModel): string {
+        const price = model.cost.input === 0 ? 'free' : `$${model.cost.input}/M in`
+        return `${Math.round(model.contextWindow / 1000)}k ctx  ${price}`
+    }
+
+    /**
+     * pi has no model picker an extension can open, so the dialog carries those columns. Sorted by
+     * price, since the whole point of the executor is what it costs.
+     */
+    function modelChoices(ctx: ExtensionContext): { label: string; ref: string }[] {
+        return physicalModels(ctx)
+            .toSorted((a, b) => a.cost.input - b.cost.input)
+            .map((model) => ({ label: `${modelRef(model)}  ${priceLine(model)}`, ref: modelRef(model) }))
+    }
+
+    pi.registerVirtualModel({
+        provider: ROUTER_PROVIDER,
+        id: ROUTER_ID,
+        name: 'Prewalk',
+        thinkingLevels: ['off', 'low', 'medium', 'high'],
+        async route(request, ctx) {
+            const level = request.thinkingLevel
+            const inPlay = request.previous
+
+            if (request.reason === 'retry' && request.failed) {
+                return { model: request.failed.model, thinkingLevel: request.failed.thinkingLevel ?? level }
+            }
+
+            const executor = ensureState(ctx)?.swapped === true ? executorModel(ctx) : null
+            if (request.reason === 'continuation') {
+                const escalated = await escalationRoute(ctx, inPlay, level)
+                if (escalated) return escalated
+            }
+
+            if (state?.reviewing === true && (prompt === null || request.reason !== 'user')) {
+                const reviewer = chainModel(ctx, config.planner)
+                if (reviewer) return { model: reviewer, thinkingLevel: level }
+            }
+            if (prompt === null || request.reason !== 'user') {
+                const asked = await steerRoute(ctx, executor, level)
+                if (asked) return asked
+                if (request.reason !== 'user' && executor) return { model: executor, thinkingLevel: level }
+                if (inPlay) return { model: inPlay.model, thinkingLevel: inPlay.thinkingLevel ?? level }
+                return { model: coldStart(ctx), thinkingLevel: level }
+            }
+            prompt = null
+
+            const decided = decideRoute()
+            startTask(ctx, decided.route)
+            if (!state) return { model: coldStart(ctx), thinkingLevel: level }
+            if (decided.route === 'run') {
+                const run = executorModel(ctx)
+                if (!run) {
+                    return heldModel(
+                        ctx,
+                        inPlay,
+                        'prewalk: no executor in the chain resolves, or has an API key',
+                        level
+                    )
+                }
+                return { model: run, thinkingLevel: level }
+            }
+            const planned = chainModel(ctx, config.planner)
+            if (!planned) {
+                return heldModel(ctx, inPlay, 'prewalk: no planner in the chain resolves, or has an API key', level)
+            }
+            state = { ...state, planner: modelRef(planned) }
+            persist()
+            return { model: planned, thinkingLevel: level }
+        }
+    })
+
+    pi.on('before_agent_start', async () => setPlanTool())
+    pi.on('turn_start', async () => {
+        demandedThisTurn = false
+    })
+    pi.on('session_start', async (_event, ctx) => reconstructState(ctx))
+    pi.on('session_tree', async (_event, ctx) => reconstructState(ctx))
+    pi.on('model_select', async (event, ctx) => claimRouter(ctx, event.model))
+
+    pi.on('agent_end', async (_event, ctx) => {
+        if (!state || state.swapped || !state.plan || state.stepIndex > 0) return
+        markSwapped(ctx)
+        stallNudged = true
+        const first = state.plan.steps[0]
+        const nudge = first
+            ? `${NUDGE_PREFIX} Begin step 1: ${first.file}: ${first.change}. ${first.verify ? `Run \`${first.verify}\` to verify.` : ''}`
+            : `${NUDGE_PREFIX} Begin the plan.`
+        pi.sendMessage(
+            { customType: 'prewalk-stall-nudge', content: nudge, display: false, details: undefined },
+            { deliverAs: 'steer' }
+        )
+    })
+
+    pi.on('agent_settled', async (event, ctx) => {
+        if (!state) return
+        if (state.inFlight) {
+            state = { ...state, inFlight: false }
+            persist()
+            setPlanTool()
+            updateStatus(ctx)
+        }
+        if (event.aborted) return
+        if (!state.swapped || !state.plan || ctx.hasPendingMessages()) return
+        const done = state.stepIndex >= state.plan.steps.length
+        if (config.review && !reviewSpent && (done || stallNudged)) {
+            startReview(ctx)
+            return
+        }
+        if (completionNudged || stallNudged) return
+        if (done) completionNudged = true
+        stallNudged = true
+        pi.sendUserMessage(continuationNudge(state.plan, state.stepIndex), { deliverAs: 'steer' })
+    })
+
+    pi.on('input', async (event) => {
+        const text = event.text.trimStart()
+        if (text.startsWith('/') || text.startsWith(SKILL_BLOCK_PREFIX) || text.startsWith(NUDGE_PREFIX)) return
+        stallNudged = false
+        // A steer joins the task in flight and a follow-up waits for that turn to settle, so the
+        // steer never opens a task here and the follow-up routes as one when it arrives. Every
+        // other message opens a scored task, whichever extension sent it.
+        if (event.streamingBehavior === 'steer') {
+            steer = event.text
+            return
+        }
+        recordPrompt(event.text)
+    })
+
+    pi.on('tool_call', async (event, ctx) => {
+        if (!state) return
+        const input = isToolInputSchema(event.input) ? event.input : undefined
+        const inert =
+            event.toolName.toLowerCase() === 'bash'
+                ? !bashWritesFiles(input?.command ?? '', ctx.cwd)
+                : isScratchPath(input?.path ?? input?.file ?? '', ctx.cwd)
+        const decision = swapDecision(state, event.toolName, inert, demandedThisTurn)
+        if (decision === 'demand-plan') {
+            demandedThisTurn = true
+            state = { ...state, planDemanded: true }
+            persist()
+            return { block: true, reason: PLAN_DEMAND_REASON }
+        }
+        if (decision === 'swap') {
+            markSwapped(ctx)
+            return
+        }
+        if (postPlanCalls === null || state.swapped || state.reviewing || event.toolName === PLAN_TOOL_NAME) return
+        postPlanCalls += 1
+        if (postPlanCalls >= PLANNER_GRACE_CALLS) markSwapped(ctx)
+    })
+
+    pi.on('tool_call', async (event) => {
+        if (!state || event.toolName !== PLAN_TOOL_NAME || state.swapped) return
+        if (!isPlanParameters(event.input)) return
+        const plan = event.input
+        stallNudged = false
+        completionNudged = false
+        reviewSpent = false
+        postPlanCalls = 0
+        state = { ...state, plan, planDemanded: true, reviewing: false, stepIndex: 0, failures: 0, wall: null }
+        persist()
+    })
+
+    pi.on('tool_result', async (event, ctx) => {
+        if (!state?.swapped) return
+        const plan = state.plan
+        const blocks = Array.isArray(event.content)
+            ? event.content.map((b) => (b.type === 'text' ? (b.text ?? '') : ''))
+            : [String(event.content ?? '')]
+        const last = blocks[blocks.length - 1] ?? ''
+        if (event.isError) {
+            const failures = state.failures + 1
+            const wall = failures >= ESCALATION_BAR.failures ? last.trim().slice(0, MAX_WALL_CHARS) : state.wall
+            const changed = wall !== state.wall
+            state = { ...state, failures, wall }
+            if (changed) persist()
+            if (plan === null || failures !== ESCALATION_BAR.failures) return
+            if (state.escalations >= ESCALATION_BAR.upgrades || config.mode !== 'plan') return
+            blocks[blocks.length - 1] = `${last}\n${wallBrief(plan, state.stepIndex, failures)}`
+            return { content: [{ type: 'text', text: blocks.join('\n') }] }
+        }
+        if (state.failures > 0 || state.wall !== null) state = { ...state, failures: 0, wall: null }
+        if (plan === null) return
+        if (!isToolInputSchema(event.input)) return
+        const tool = event.toolName.toLowerCase()
+        const path = event.input.path ?? event.input.file ?? ''
+        const command = event.input.command ?? ''
+        const matched = READ_ONLY_TOOLS.has(tool)
+            ? -1
+            : path === ''
+              ? stepForCommand(plan, state.stepIndex, command)
+              : stepForPath(plan, state.stepIndex, path)
+        if (matched === -1) return
+        const { index, footer } = planProgress(plan, state.stepIndex, matched)
+        if (index !== state.stepIndex) {
+            state = { ...state, stepIndex: index, failures: 0, wall: null }
+            stallNudged = false
+            persist()
+            updateStatus(ctx)
+        }
+        blocks[blocks.length - 1] = `${last}\n${footer}`
+        return { content: [{ type: 'text', text: blocks.join('\n') }] }
+    })
+
+    pi.on('context', async (event) => {
+        if (!state) return
+        return { messages: contextForModel(state, event.messages) }
+    })
+
+    const planTool: PlanToolDefinition = {
+        name: PLAN_TOOL_NAME,
+        label: 'Commit Plan',
+        defaultActive: false,
+        description: PLANNING_BRIEF,
+        promptSnippet: 'Commit the implementation plan before your first edit',
+        annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: true, openWorldHint: false },
+        exposure: 'deferred',
+        namespace: {
+            name: 'prewalk',
+            description:
+                'Commit the implementation plan before your first edit, when a prewalk has the session planning.'
+        },
+        promptGuidelines: [
+            `While a prewalk has you planning, search for ${PLAN_TOOL_NAME} and call it exactly once, after you finish exploring and before your first edit.`,
+            'The executor will not re-read what you read. Each step must carry enough detail to land the edit cold: the file, the exact anchor text to find, what changes, and the verify command.',
+            'If a step requires a judgement the executor cannot make, resolve it now or split the step.'
+        ],
+        parameters: PLAN_PARAMETERS,
+        async execute(_toolCallId, params: PlanParams, _signal, _onUpdate, _ctx) {
+            const text = renderHandoff({
+                summary: params.summary,
+                facts: params.facts,
+                steps: params.steps,
+                check: params.check,
+                outOfScope: params.outOfScope
+            })
+            return {
+                content: [{ type: 'text', text }],
+                details: { summary: params.summary, steps: params.steps.length }
+            }
+        }
+    }
+    pi.registerTool(planTool)
+
+    pi.registerCommand('prewalk', {
+        description: 'Route each task to a planning model or a cheap one, by the phase or by hand',
+        getArgumentCompletions: (prefix: string) => {
+            const lastSpace = prefix.lastIndexOf(' ')
+            const head = lastSpace === -1 ? '' : prefix.slice(0, lastSpace + 1)
+            const p = prefix.slice(lastSpace + 1).toLowerCase()
+            const parsedHead = parseArgs(head, makeSpecPredicate(models))
+            if (parsedHead.kind !== 'bare' && parsedHead.kind !== 'start') return null
+            if (parsedHead.kind === 'start' && (parsedHead.task !== '' || parsedHead.planner)) return null
+            const items: { value: string; label: string; description?: string }[] = []
+            if (head === '') {
+                if (state?.executor) {
+                    items.push({
+                        value: state.executor,
+                        label: `${shortName(state.executor)} (last used)`,
+                        description: 're-arm with this executor'
+                    })
+                }
+                for (const [value, description] of [
+                    ['models', 'both chains, what resolves, what it costs'],
+                    ['mode plan', 'plan each task, no routing call'],
+                    ['mode off', 'the cheap model, no plan and no call'],
+                    ['planner', 'set the planner chain'],
+                    ['executor', 'set the executor chain'],
+                    ['plan', 'next task plans, by hand'],
+                    ['run', 'next task runs on the cheap model'],
+                    ['off', 'disarm prewalk']
+                ] as const) {
+                    items.push({ value, label: value, description })
+                }
+            }
+            const slot = parsedHead.kind === 'start' && parsedHead.executor ? 'executor' : 'planner or executor'
+            for (const m of models) {
+                items.push({ value: modelRef(m), label: m.id, description: `${m.provider} · ${slot}` })
+            }
+            const filtered = items.filter(
+                (item) => item.value.toLowerCase().includes(p) || item.label.toLowerCase().includes(p)
+            )
+            return filtered.length > 0 ? filtered.slice(0, 30).map((i) => ({ ...i, value: head + i.value })) : null
+        },
+        handler: async (args, ctx) => {
+            const parsed = parseArgs(args, makeSpecPredicate(physicalModels(ctx)))
+
+            if (parsed.kind === 'status') {
+                const report = configReport(ctx)
+                ctx.ui.notify(state ? `${statusLabel(state)}\n${report}` : `prewalk is not armed.\n${report}`, 'info')
+                return
+            }
+
+            if (parsed.kind === 'stop') {
+                if (!state) {
+                    ctx.ui.notify('No prewalk active', 'info')
+                    return
+                }
+                disarm(ctx)
+                ctx.ui.notify('Prewalk disarmed', 'info')
+                return
+            }
+
+            if (parsed.kind === 'models') {
+                ctx.ui.notify(configReport(ctx), 'info')
+                return
+            }
+
+            if (parsed.kind === 'mode') {
+                config = { ...config, mode: parsed.mode }
+                writePrewalkConfig(config)
+                ctx.ui.notify(`prewalk: ${MODE_LABEL[parsed.mode]}`, 'info')
+                return
+            }
+
+            if (parsed.kind === 'role') {
+                if (parsed.chain.length === 0) {
+                    const choices = modelChoices(ctx)
+                    const picked =
+                        choices.length === 0
+                            ? undefined
+                            : await ctx.ui.select(
+                                  `prewalk: pick the ${parsed.role} model`,
+                                  choices.map((choice) => choice.label)
+                              )
+                    const choice = choices.find((candidate) => candidate.label === picked)
+                    if (choice === undefined) return
+                    parsed.chain.push(choice.ref)
+                }
+                config = { ...config, [parsed.role]: parsed.chain }
+                writePrewalkConfig(config)
+                // The badge reads the armed pair, so a chain that moves under a live session moves it too.
+                if (state !== null) {
+                    const head = chainModel(ctx, parsed.chain)
+                    state = { ...state, [parsed.role]: head ? modelRef(head) : parsed.chain[0] }
+                    persist()
+                    updateStatus(ctx)
+                }
+                ctx.ui.notify(`prewalk: ${parsed.role} chain is ${parsed.chain.map(shortName).join(' -> ')}`, 'info')
+                return
+            }
+
+            if (parsed.kind === 'force') {
+                await arm(ctx, { planner: null, executor: null, task: parsed.task, forced: parsed.route })
+                return
+            }
+
+            if (parsed.kind === 'bare') {
+                if (state) {
+                    ctx.ui.notify(statusLabel(state), 'info')
+                    return
+                }
+                await arm(ctx, { planner: null, executor: null, task: '' })
+                return
+            }
+            await arm(ctx, parsed)
+        }
+    })
+}