diff --git a/docs/plan-approval.md b/docs/plan-approval.md index 4aabb7be..5fc27e61 100644 --- a/docs/plan-approval.md +++ b/docs/plan-approval.md @@ -2,9 +2,14 @@ `exit_plan_mode` submits a written proposal for review. Leaving plan mode and approving the proposal are separate events. In headless and RPC sessions, a valid -proposal exits plan mode and restores the previous tools, while the caller remains -responsible for obtaining approval before execution. This also applies to RPC -contexts with `hasUI: true`; their UI bridge does not provide interactive approval. +proposal exits plan mode, while the caller remains responsible for obtaining +approval before execution. This also applies to RPC contexts with `hasUI: true`; +their UI bridge does not provide interactive approval. + +Plan mode does not hide tools: it denies mutating calls outside the plan file +through the `tool_call` policy, so the request's tool block stays byte-identical +and the provider's prompt cache keeps its prefix. Only a profile that lacks a +plan tool has it appended on entry, and that addition alone is reverted on exit. Successful headless and RPC exits return the existing model-facing text plus structured `details`, for example: diff --git a/packages/coding-agent/docs/step-integration.md b/packages/coding-agent/docs/step-integration.md index 527b8702..7402ed3e 100644 --- a/packages/coding-agent/docs/step-integration.md +++ b/packages/coding-agent/docs/step-integration.md @@ -328,10 +328,14 @@ for path rules and retention responsibilities. effect immediately; it does not request approval. An explicit `--plan` applies at process startup, including a continued session whose mode was off; later reloads and session/branch navigation do not reapply it. The plan path is - announced, and `write_file` / `edit_file` remain available for that file. - Other file-write targets are blocked while planning. This is **not a shell - sandbox**: command tools retain the existing permission policy, and the model - is instructed to keep commands read-only while planning. + announced. Plan mode does not remove tools from the request — that would + invalidate the provider's cached prompt prefix — so enforcement lives in the + `tool_call` policy: `write_file`, `edit_file`, `write`, and `edit` may only + target the plan file, and every other target is rejected with an actionable + denial. Other file-writing tools, including MCP servers, are not covered by + that check. This is **not a shell sandbox**: command tools retain the existing + permission policy, and the model is instructed to keep commands read-only + while planning. - `exit_plan_mode` submits the written proposal for review, requiring a readable, nonempty regular plan file. In the TUI, the user can approve execution, stay in plan mode, or provide refinement notes. Cancel @@ -415,8 +419,11 @@ The existing six tool names and `/todos` remain; the task schemas add explicit new-plan, history-query, and resume fields. Legacy support remains limited to migration of legacy plan-mode todos into tasks on the active branch and restoring existing task snapshots as a single plan. Older -saved tool lists retain currently active plan/task tools on restore and exit; -ordinary custom tools still follow the saved branch selection. There is no +saved tool lists — written by builds that narrowed the active list on entry — +retain currently active plan/task tools on restore and exit; ordinary custom +tools still follow the saved branch selection. Builds that append only the plan +tools a profile lacks (`version: 2` entries) store no replaced list, so a branch +written by them never replays an older snapshot. There is no additional `TodoWrite` / `update_plan` alias or second checklist store. The separation follows [Claude Code's plan and task tools](https://code.claude.com/docs/en/tools-reference) diff --git a/packages/coding-agent/src/features/plan-mode-tools.ts b/packages/coding-agent/src/features/plan-mode-tools.ts index d8e46fad..e7e22b57 100644 --- a/packages/coding-agent/src/features/plan-mode-tools.ts +++ b/packages/coding-agent/src/features/plan-mode-tools.ts @@ -32,7 +32,7 @@ export interface StepPlanModeController { * /plan toggle owned by the extension. Returns the session plan file path. */ enterPlanMode(extensionContext: ExtensionContext): string; - /** Leave plan mode: restore tools, persist state, record telemetry. */ + /** Leave plan mode: restore any tools plan mode added, persist state, record telemetry. */ exitPlanMode(outcome: StepPlanExitOutcome, extensionContext: ExtensionContext): void; /** Plan file path for this session (persisted or derived from the session id). */ resolvePlanFilePath(extensionContext: ExtensionContext): string; diff --git a/packages/coding-agent/src/features/step-plan.ts b/packages/coding-agent/src/features/step-plan.ts index ebf178ec..62e01742 100644 --- a/packages/coding-agent/src/features/step-plan.ts +++ b/packages/coding-agent/src/features/step-plan.ts @@ -7,10 +7,10 @@ */ import { existsSync } from "node:fs"; -import path from "node:path"; import type { AgentMessage } from "@step-harness/agent-core"; import type { EventBus } from "../core/event-bus.ts"; import type { ExtensionAPI, ExtensionContext, ExtensionFactory } from "../core/extensions/types.ts"; +import { resolveToCwd } from "../core/tools/path-utils.ts"; import { type StepTelemetryReporter, trackStepTelemetry } from "../step/telemetry.ts"; import { type LegacyPlanTodoFields, migrateLegacyPlanTodos } from "./plan-mode-migration.ts"; import { @@ -44,7 +44,15 @@ const PLAN_TASK_TOOLS = new Set([ export type StepPlanSource = "user" | "agent" | "unknown"; interface PlanState { + /** Written by this build. Its presence marks an entry that never narrowed the tool list. */ + version?: 2; enabled: boolean; + /** + * Legacy only: the list a narrowing entry replaced. Older builds dropped the + * editing tools on entry and could only restore by replaying this snapshot, + * so sessions written by them still rely on it. This build appends only the + * plan tools a profile lacks and remembers just those additions. + */ toolsBeforePlanMode?: string[]; planFilePath?: string; /** "user" for /plan or --plan, "agent" for enter_plan_mode. */ @@ -55,6 +63,11 @@ function unique(values: readonly string[]): string[] { return [...new Set(values)]; } +/** True when both lists hold the same tool names in the same order. */ +function sameToolOrder(left: readonly string[], right: readonly string[]): boolean { + return left.length === right.length && left.every((name, index) => name === right[index]); +} + const publishedPlanModes = new WeakMap(); export function getStepPlanMode(manager: object): string { const value = publishedPlanModes.get(manager); @@ -102,24 +115,63 @@ export const createStepPlanExtension = (options: { telemetry?: StepTelemetryReporter } = {}): ExtensionFactory => (pi: ExtensionAPI): void => { let enabled = false; - let toolsBeforePlanMode: string[] | undefined; + /** Plan tools this process appended on entry; the only thing exit removes. */ + let planModeAddedTools: string[] | undefined; let planFilePath: string | undefined; let planSource: StepPlanSource | undefined; // The event bus is optional so minimal embedder/test harnesses that stub // ExtensionAPI keep working without one. const events = (pi as { events?: EventBus }).events; - const persistPlanState = (): void => { - pi.appendEntry("step-plan", { enabled, toolsBeforePlanMode, planFilePath, planSource }); + /** + * Persist the branch's plan state. + * + * `carriesLegacyTools` is only set when rewriting a legacy entry during + * migration: the entry it replaces is the sole record of the list plan mode + * narrowed, so dropping it would strand that session on the next restore. + * Ordinary entries never carry a snapshot — this build appends tools + * instead of replacing the list. + */ + const persistPlanState = (carriesLegacyTools?: string[]): void => { + pi.appendEntry("step-plan", { + version: 2, + enabled, + planFilePath, + planSource, + ...(carriesLegacyTools ? { toolsBeforePlanMode: carriesLegacyTools } : {}), + }); }; + /** + * Plan mode is a permission narrowing, not a capability change. + * + * The model keeps every tool it already had; the `tool_call` hook below + * denies mutations outside the plan file. Rebuilding the list would rewrite + * the request's tool block, which the provider serializes ahead of the + * system prompt and messages: one such edit invalidates the cached prefix + * for the whole conversation, so a single plan entry can re-bill hundreds of + * thousands of prompt tokens. + * + * A session that does not already offer every plan tool still has to gain + * them (a Pi-native profile lacks the Step aliases). Append only what is + * missing and keep the existing order, so the common Step profile — which + * already offers all of them — leaves the tool block byte-identical. + */ const enable = (): void => { - toolsBeforePlanMode ??= pi.getActiveTools(); - // Keep plan-file writes available; tool_call restricts their target. - pi.setActiveTools(unique([...toolsBeforePlanMode.filter((name) => !MUTATING_TOOLS.has(name)), ...PLAN_TOOLS])); + const active = pi.getActiveTools(); + const missing = PLAN_TOOLS.filter((name) => !active.includes(name)); + if (missing.length === 0) return; + planModeAddedTools = missing; + pi.setActiveTools([...active, ...missing]); }; const restoreTools = (): void => { - if (toolsBeforePlanMode) pi.setActiveTools(toolsBeforePlanMode); - toolsBeforePlanMode = undefined; + const added = planModeAddedTools; + planModeAddedTools = undefined; + if (!added || added.length === 0) return; + // Remove only what entering added, so tools activated in the meantime + // (a late MCP catalog, find_tools) keep their place in the list. + const live = pi.getActiveTools(); + const remaining = live.filter((name) => !added.includes(name)); + if (remaining.length !== live.length) pi.setActiveTools(remaining); }; const resolvePlanFilePath = (ctx: ExtensionContext): string => planFilePath ?? getPlanFilePath(ctx.sessionManager.getSessionId(), ctx.cwd); @@ -185,9 +237,11 @@ export const createStepPlanExtension = pi.on("tool_call", async (event, ctx) => { if (!enabled || !MUTATING_TOOLS.has(event.toolName)) return; const planPath = resolvePlanFilePath(ctx); - const toolInput = event.input as Record; + const toolInput = (event.input ?? {}) as Record; const targetPath = typeof toolInput.path === "string" ? toolInput.path : ""; - if (targetPath && path.resolve(ctx.cwd, targetPath) === path.resolve(planPath)) { + // Resolve exactly as the file tools do (`~`, `@`, unicode spaces), so an + // allowed target cannot land somewhere else than the guard checked. + if (targetPath && resolveToCwd(targetPath, ctx.cwd) === resolveToCwd(planPath, ctx.cwd)) { // Allowed: the mutation targets the plan file itself. This runs // before the write lands, so a missing file marks the plan's creation. recordPlanUpdate(options.telemetry, planSource, !existsSync(planPath)); @@ -195,7 +249,13 @@ export const createStepPlanExtension = } return { block: true, - reason: `Plan mode blocked this file mutation. Only the plan file may be written: ${planPath}\nTarget: ${targetPath || "(no path provided)"}`, + reason: [ + `Plan mode is active, so this file mutation was blocked. Only the plan file may be written: ${planPath}`, + `Target: ${targetPath || "(no path provided)"}`, + "Plan mode is a permission restriction, not a missing tool: do not retry this edit or rephrase the path.", + "run_command keeps its normal permissions, but stay read-only while planning: do not use it to write or patch files.", + `Write the implementation proposal to the plan file with write_file, then call exit_plan_mode to submit it for review.`, + ].join("\n"), }; }); @@ -218,7 +278,7 @@ export const createStepPlanExtension = // Legacy tool lists predate these capabilities. Preserve only the ones // already active; ordinary tools still follow the saved branch selection. const activePlanTaskTools = pi.getActiveTools().filter((name) => PLAN_TASK_TOOLS.has(name)); - // Undo the previous branch's tool restriction before reading the new one. + // Undo whatever the previous branch's entry applied before reading the next. restoreTools(); const planStates = [...ctx.sessionManager.getBranch()] .reverse() @@ -228,24 +288,30 @@ export const createStepPlanExtension = : [], ); const restoredPlanState = planStates[0]; - // Off snapshots omit transient tools; a preceding plan entry retains the baseline. - const savedTools = planStates.find((state) => state?.toolsBeforePlanMode)?.toolsBeforePlanMode; + // The list to undo belongs to whichever entry narrowed it. An on-entry + // wins over an older carried snapshot, so a newer plan session never + // replays an earlier one's list; an off branch falls back to the last + // entry that recorded one. This build records nothing on entry, so its + // own entries only reach here after carrying a legacy list forward. + const governingEntry = planStates.find((state) => state?.enabled || state?.toolsBeforePlanMode !== undefined); + const legacyBaseline = governingEntry?.toolsBeforePlanMode; enabled = restoredPlanState?.enabled ?? false; - toolsBeforePlanMode = savedTools ? unique([...savedTools, ...activePlanTaskTools]) : undefined; planFilePath = restoredPlanState?.planFilePath; planSource = enabled ? (restoredPlanState?.planSource ?? "unknown") : undefined; + if (legacyBaseline) { + const baseline = unique([...legacyBaseline, ...activePlanTaskTools]); + if (!sameToolOrder(pi.getActiveTools(), baseline)) pi.setActiveTools(baseline); + } if (enabled) { planFilePath = resolvePlanFilePath(ctx); enable(); - } else { - restoreTools(); } if (restoredPlanState) { migrateLegacyPlanTodos({ persistedPlanState: restoredPlanState, events, extensionContext: ctx, - persistMigratedPlanState: persistPlanState, + persistMigratedPlanState: () => persistPlanState(legacyBaseline), }); } if (startWithPlan && !enabled) { diff --git a/packages/coding-agent/test/step-plan-extension.test.ts b/packages/coding-agent/test/step-plan-extension.test.ts index 73ef5ee3..d584f188 100644 --- a/packages/coding-agent/test/step-plan-extension.test.ts +++ b/packages/coding-agent/test/step-plan-extension.test.ts @@ -41,6 +41,18 @@ interface Harness { ctx: (overrides?: Record) => ExtensionContext; } +/** Every tool the plan-mode entry has to guarantee, in Step's model-facing names. */ +const PLAN_TOOLS_UNDER_TEST = [ + "read_file", + "find_files", + "search_files", + "list_directory", + "run_command", + "clarify_user", + "write_file", + "edit_file", +]; + /** Keys the review dialog reacts to, spelled the way the terminal sends them. */ const REVIEW_KEYS = { enter: "\r", escape: "\u001b", down: "\u001b[B", up: "\u001b[A" } as const; @@ -656,25 +668,52 @@ test.each(planEntryCases)( expect(fs.existsSync(planDir)).toBe(false); const source = entry === "enter_plan_mode" ? "agent" : "user"; expect(lastPlanEntry(harness)).toMatchObject({ + version: 2, enabled: true, planSource: source, planFilePath: planPath, - toolsBeforePlanMode: originalTools, }); + // This build appends instead of narrowing, so it stores no replaced list. + expect(lastPlanEntry(harness)).not.toHaveProperty("toolsBeforePlanMode"); expect(harness.activeTools()).toEqual( expect.arrayContaining(["write_file", "edit_file", "run_command", "custom_tool"]), ); - expect(harness.activeTools()).not.toContain("write"); - expect(harness.activeTools()).not.toContain("edit"); + // Plan mode narrows permissions, not the tool block: the aliases stay + // visible and the tool_call policy below is what blocks them. Dropping + // them would rewrite the request's tool block and discard the provider's + // cached prompt prefix for the whole conversation. + expect(harness.activeTools()).toContain("write"); + expect(harness.activeTools()).toContain("edit"); if (source === "user") { expect(harness.notifications.at(-1)?.message).toContain(planPath); expect(harness.notifications.at(-1)?.message).not.toContain("command execution are disabled"); } + // The guard is the only enforcement now that the editing tools stay + // visible, and the file tools resolve through resolveToCwd. Every accepted + // spelling below must therefore mean the same file the tools will write. + const accepted = [ + path.relative(cwd, planPath), + planPath, + `@${path.relative(cwd, planPath)}`, + path.join(cwd, "..", path.basename(cwd), path.relative(cwd, planPath)), + ]; + // A narrowed-but-not-normalized target stays blocked: these all point + // somewhere else even though they look like the plan path. + const rejected = [ + "src/app.ts", + ".stepcode/config.toml", + path.join(planDir, "other-plan.md"), + "@src/app.ts", + " src/app.ts", + path.join(path.relative(cwd, planPath), "..", "..", "evil.md"), + ]; for (const toolName of ["write_file", "edit_file", "write", "edit"]) { - await expect( - harness.emit({ type: "tool_call", toolName, input: { path: path.relative(cwd, planPath) } }, context), - ).resolves.toBeUndefined(); - for (const target of ["src/app.ts", ".stepcode/config.toml", path.join(planDir, "other-plan.md")]) { + for (const target of accepted) { + await expect( + harness.emit({ type: "tool_call", toolName, input: { path: target } }, context), + ).resolves.toBeUndefined(); + } + for (const target of rejected) { await expect( harness.emit({ type: "tool_call", toolName, input: { path: target } }, context), ).resolves.toMatchObject({ block: true }); @@ -706,6 +745,104 @@ test.each(planEntryCases)( }, ); +test("entering and leaving plan mode keeps an equivalent tool set in its original order", async () => { + const cwd = makeWorkspace(); + const harness = createHarness({ cwd }); + // A real Step profile already offers every plan tool, including clarify_user, + // so plan mode narrows permissions without changing the capability set. + const originalTools = [ + "clarify_user", + "enter_plan_mode", + "exit_plan_mode", + "list_directory", + "find_files", + "search_files", + "read_file", + "write_file", + "edit_file", + "run_command", + "find_tools", + "custom_tool", + ]; + harness.api.setActiveTools(originalTools); + const applied: string[][] = []; + const setActiveTools = harness.api.setActiveTools.bind(harness.api); + harness.api.setActiveTools = (names: string[]) => { + applied.push([...names]); + setActiveTools(names); + }; + + createStepPlanExtension()(harness.api); + const context = harness.ctx(); + await harness.commands.get("plan")!.handler("", context); + + // Reapplying the same set in a different order rewrites the request's tool + // block and drops the provider prompt cache for the whole conversation. + expect(harness.activeTools()).toEqual(originalTools); + expect(applied).toEqual([]); + + // Plan-file writes stay gated by the tool_call path check, not by absence. + const planPath = path.join(cwd, ".stepcode", "plans", "session-sess-1.md"); + await expect( + harness.emit( + { type: "tool_call", toolName: "write_file", input: { path: path.relative(cwd, planPath) } }, + context, + ), + ).resolves.toBeUndefined(); + await expect( + harness.emit({ type: "tool_call", toolName: "write_file", input: { path: "src/app.ts" } }, context), + ).resolves.toMatchObject({ block: true }); + + await harness.commands.get("plan")!.handler("", context); + expect(harness.activeTools()).toEqual(originalTools); + expect(applied).toEqual([]); +}); + +test("entering plan mode appends only the plan tools a profile lacks", async () => { + const cwd = makeWorkspace(); + const harness = createHarness({ cwd }); + // A Pi-native profile has the editing aliases but not the Step names, so + // plan mode has to add them — and must not remove anything to do it. + const originalTools = ["read_file", "write", "edit", "custom_tool"]; + harness.api.setActiveTools(originalTools); + const applied: string[][] = []; + const setActiveTools = harness.api.setActiveTools.bind(harness.api); + harness.api.setActiveTools = (names: string[]) => { + applied.push([...names]); + setActiveTools(names); + }; + + createStepPlanExtension()(harness.api); + const context = harness.ctx(); + + await harness.commands.get("plan")!.handler("", context); + expect(applied).toEqual([ + [ + ...originalTools, + "find_files", + "search_files", + "list_directory", + "run_command", + "clarify_user", + "write_file", + "edit_file", + ], + ]); + // Existing order is preserved and nothing is dropped. + expect(harness.activeTools().slice(0, originalTools.length)).toEqual(originalTools); + expect(harness.activeTools()).toContain("write"); + expect(harness.activeTools()).toContain("edit"); + + // The aliases stay visible; the policy blocks the mutation instead. + await expect( + harness.emit({ type: "tool_call", toolName: "write", input: { path: "src/app.ts" } }, context), + ).resolves.toMatchObject({ block: true }); + + await harness.commands.get("plan")!.handler("", context); + expect(harness.activeTools()).toEqual(originalTools); + expect(applied).toHaveLength(2); +}); + async function restoreLifecycle(harness: Harness, type: "session_start" | "session_tree"): Promise { await harness.emit( type === "session_start" @@ -817,14 +954,155 @@ test("legacy off snapshots recover the latest branch baseline and new entry capt expect(harness.appended).toEqual([]); // Restoring an off snapshot must not leave a transient planning baseline. + // This build appends instead of narrowing, so a new entry records no replaced + // list and exit reverses only what the entry actually added. const changedTools = ["read_file", "new_tool"]; harness.api.setActiveTools(changedTools); await harness.commands.get("plan")!.handler("", harness.ctx()); - expect(lastPlanEntry(harness).toolsBeforePlanMode).toEqual(changedTools); + expect(lastPlanEntry(harness)).toMatchObject({ version: 2, enabled: true }); + expect(lastPlanEntry(harness)).not.toHaveProperty("toolsBeforePlanMode"); await harness.commands.get("plan")!.handler("", harness.ctx()); expect(harness.activeTools()).toEqual(changedTools); }); +test("exiting plan mode removes only the tools entering added", async () => { + const harness = createHarness(); + // A Pi-native profile lacks the Step aliases, so entering really does add. + harness.api.setActiveTools(["read_file", "write", "edit", "custom_tool"]); + createStepPlanExtension()(harness.api); + // Register the custom names too: the real session drops unknown names, so a + // test that relies on them must not depend on the stub echoing everything. + for (const name of ["custom_tool", "late_tool"]) { + harness.api.registerTool({ ...harness.tools.get("enter_plan_mode")!, name }); + } + const context = harness.ctx(); + + await harness.commands.get("plan")!.handler("", context); + expect(harness.activeTools()).toEqual( + expect.arrayContaining(["read_file", "write", "edit", "custom_tool", "write_file", "edit_file"]), + ); + + // A tool activated while planning (a late MCP catalog, find_tools) must keep + // its place: exit undoes its own addition instead of replaying a snapshot. + harness.api.setActiveTools([...harness.activeTools(), "late_tool"]); + await harness.commands.get("plan")!.handler("", context); + + expect(harness.activeTools()).toEqual(["read_file", "write", "edit", "custom_tool", "late_tool"]); +}); + +test("a legacy todo migration keeps the replaced tool list for later restores", async () => { + // Migration rewrites the entry. Dropping the snapshot here would strand the + // branch: the next restore could no longer tell which list plan mode replaced. + const originalTools = ["read_file", "write", "edit"]; + const harness = createHarness({ + entries: [ + { + type: "custom", + customType: "step-plan", + data: { + enabled: true, + planSource: "user", + toolsBeforePlanMode: originalTools, + todos: [{ text: "Legacy task" }], + }, + }, + ], + }); + createStepPlanExtension()(harness.api); + await restoreLifecycle(harness, "session_start"); + + const migrated = lastPlanEntry(harness); + expect(migrated).toMatchObject({ version: 2, enabled: true }); + expect(migrated.toolsBeforePlanMode).toEqual(originalTools); + + // A fresh process must replay that list, not fall back to the profile default. + harness.session.resetLeaf?.(); + await restoreLifecycle(harness, "session_start"); + expect(harness.activeTools()).toEqual(expect.arrayContaining(originalTools)); +}); + +test("a plan entry that adds nothing records no replaced tool list", async () => { + const harness = createHarness(); + const live = [...PLAN_TOOLS_UNDER_TEST, "find_tools"]; + harness.api.setActiveTools(live); + const applied: string[][] = []; + const setActiveTools = harness.api.setActiveTools.bind(harness.api); + harness.api.setActiveTools = (names: string[]) => { + applied.push([...names]); + setActiveTools(names); + }; + createStepPlanExtension()(harness.api); + + await harness.commands.get("plan")!.handler("", harness.ctx()); + expect(lastPlanEntry(harness)).not.toHaveProperty("toolsBeforePlanMode"); + await harness.commands.get("plan")!.handler("", harness.ctx()); + + // Enter, exit, and the persisted snapshot all leave the list alone. + expect(applied).toEqual([]); + expect(harness.activeTools()).toEqual(live); +}); + +test("restoring a new-format plan branch reapplies nothing", async () => { + const harness = createHarness({ + entries: [{ type: "custom", customType: "step-plan", data: { version: 2, enabled: true } }], + }); + const live = [...PLAN_TOOLS_UNDER_TEST, "find_tools", "custom_tool"]; + harness.api.setActiveTools(live); + const applied: string[][] = []; + const setActiveTools = harness.api.setActiveTools.bind(harness.api); + harness.api.setActiveTools = (names: string[]) => { + applied.push([...names]); + setActiveTools(names); + }; + createStepPlanExtension()(harness.api); + + await restoreLifecycle(harness, "session_start"); + + expect(harness.statuses.get("plan-mode")).toBe("Planning (unknown)"); + expect(applied).toEqual([]); + expect(harness.activeTools()).toEqual(live); +}); + +test("a stale legacy baseline is not replayed onto a newer plan entry", async () => { + // Entry 1 narrowed the list and recorded what it replaced. Entry 3 was written + // by this build, which never narrows; replaying entry 1 here would reorder the + // current list and drop tools the newer branch never touched. + const harness = createHarness({ + entries: [ + { type: "custom", customType: "step-plan", data: { enabled: true, toolsBeforePlanMode: ["read_file"] } }, + { type: "custom", customType: "step-plan", data: { enabled: false } }, + { type: "custom", customType: "step-plan", data: { version: 2, enabled: true } }, + ], + }); + createStepPlanExtension()(harness.api); + const live = [ + "find_tools", + "read_file", + "find_files", + "search_files", + "list_directory", + "run_command", + "clarify_user", + "write_file", + "edit_file", + "late_tool", + ]; + harness.api.setActiveTools(live); + const applied: string[][] = []; + const setActiveTools = harness.api.setActiveTools.bind(harness.api); + harness.api.setActiveTools = (names: string[]) => { + applied.push([...names]); + setActiveTools(names); + }; + + await restoreLifecycle(harness, "session_start"); + expect(harness.statuses.get("plan-mode")).toBe("Planning (unknown)"); + // Every plan tool was already live, so the entry adds nothing and the stale + // baseline is never applied: same list, same order. + expect(applied).toEqual([]); + expect(harness.activeTools()).toEqual(live); +}); + test.each(["session_start", "session_tree"] as const)( "%s clears plan state on an empty branch and restores the previous tools", async (type) => { @@ -852,9 +1130,9 @@ test.each(["session_start", "session_tree"] as const)( ).resolves.toBeUndefined(); await runTool(harness, "enter_plan_mode", {}, harness.ctx()); expect(lastPlanEntry(harness)).toMatchObject({ + version: 2, planFilePath: path.join(cwd, ".stepcode", "plans", "session-sess-1.md"), planSource: "agent", - toolsBeforePlanMode: originalTools, }); }, ); diff --git a/packages/coding-agent/test/suite/regressions/plan-mode-prompt-cache-prefix.test.ts b/packages/coding-agent/test/suite/regressions/plan-mode-prompt-cache-prefix.test.ts new file mode 100644 index 00000000..b1d48425 --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/plan-mode-prompt-cache-prefix.test.ts @@ -0,0 +1,152 @@ +import { mkdirSync, writeFileSync } from "node:fs"; +import path from "node:path"; +import { fauxAssistantMessage, fauxToolCall } from "@step-harness/providers"; +import { Type } from "typebox"; +import { beforeAll, expect, it } from "vitest"; +import { getPlanFilePath } from "../../../src/features/plan-mode-tools.ts"; +import { createStepPlanExtension } from "../../../src/features/step-plan.ts"; +import { initTheme } from "../../../src/theme/theme.ts"; +import { createHarness } from "../harness.ts"; + +/** + * Plan mode is a permission narrowing, not a capability change. The provider + * serializes `tools` ahead of the system prompt and messages, so rewriting the + * tool block invalidates the cached prefix for the entire conversation. + * + * A real Step profile already offers every plan tool; entering plan mode used to + * reorder that same set and re-bill the whole prompt. Measured on a live session, + * one entry at a ~230k-token prompt read back only 1.6k tokens. + */ + +const PLAN_TOOLS = [ + "read_file", + "find_files", + "search_files", + "list_directory", + "run_command", + "clarify_user", + "write_file", + "edit_file", +] as const; + +function stepTool(name: string, execute?: (args: Record) => void) { + return { + name, + label: name, + description: `${name} tool`, + parameters: Type.Object({ path: Type.Optional(Type.String()), content: Type.Optional(Type.String()) }), + execute: async (_toolCallId: string, params: unknown) => { + execute?.((params ?? {}) as Record); + return { content: [{ type: "text" as const, text: name }], details: {} }; + }, + }; +} + +// The plan extension paints its status line through the active theme. +beforeAll(() => initTheme("dark")); + +interface CapturedRequest { + tools: string[]; + systemPrompt: string; +} + +it("keeps the request tool block and system prompt byte-identical across a plan-mode cycle", async () => { + // Filled in once the harness owns its workspace; the tool only runs later. + let workspace = ""; + const harness = await createHarness({ + tools: [...PLAN_TOOLS, "find_tools"].map((name) => + stepTool(name, (args) => { + if (name !== "write_file" || typeof args.path !== "string") return; + const target = path.resolve(workspace, args.path); + mkdirSync(path.dirname(target), { recursive: true }); + writeFileSync(target, typeof args.content === "string" ? args.content : ""); + }), + ), + extensionFactories: [{ name: "step-plan", factory: createStepPlanExtension(), hidden: true }], + }); + workspace = harness.tempDir; + + try { + // The shape of an ordinary Step profile: every plan tool is already active, + // so plan mode has nothing to add. + harness.session.setActiveToolsByName([...PLAN_TOOLS, "find_tools", "enter_plan_mode", "exit_plan_mode"]); + const planPath = getPlanFilePath(harness.sessionManager.getSessionId(), harness.tempDir); + + const captured: CapturedRequest[] = []; + const capture = (context: { tools?: Array<{ name: string }>; systemPrompt?: string }) => { + captured.push({ + tools: (context.tools ?? []).map((tool) => tool.name), + systemPrompt: context.systemPrompt ?? "", + }); + }; + harness.setResponses([ + (context) => { + capture(context); + return fauxAssistantMessage(fauxToolCall("enter_plan_mode", {}), { stopReason: "toolUse" }); + }, + (context) => { + capture(context); + return fauxAssistantMessage( + fauxToolCall("write_file", { + path: path.relative(workspace, planPath), + content: "# Plan\n\n1. Do it.\n", + }), + { stopReason: "toolUse" }, + ); + }, + (context) => { + capture(context); + return fauxAssistantMessage(fauxToolCall("exit_plan_mode", {}), { stopReason: "toolUse" }); + }, + (context) => { + capture(context); + return fauxAssistantMessage("done"); + }, + ]); + + await harness.session.prompt("Propose an approach."); + + // The system prompt embeds the active tool snippets, so equal tools plus an + // equal prompt is exactly the prefix the provider caches. + expect(captured).toHaveLength(4); + for (const request of captured.slice(1)) { + expect(request.tools).toEqual(captured[0]?.tools); + expect(request.systemPrompt).toBe(captured[0]?.systemPrompt); + } + expect(harness.getPendingResponseCount()).toBe(0); + } finally { + harness.cleanup(); + } +}); + +it("still blocks non-plan mutations after entering plan mode with tools intact", async () => { + const harness = await createHarness({ + tools: [...PLAN_TOOLS, "find_tools"].map((name) => stepTool(name)), + extensionFactories: [{ name: "step-plan", factory: createStepPlanExtension(), hidden: true }], + }); + + try { + harness.session.setActiveToolsByName([...PLAN_TOOLS, "find_tools", "enter_plan_mode", "exit_plan_mode"]); + const toolsBefore = harness.session.agent.state.tools.map((tool) => tool.name); + + harness.setResponses([ + () => fauxAssistantMessage(fauxToolCall("enter_plan_mode", {}), { stopReason: "toolUse" }), + // The editing tool is still visible; the policy must reject the target. + () => + fauxAssistantMessage(fauxToolCall("write_file", { path: "src/app.ts", content: "x" }), { + stopReason: "toolUse", + }), + () => fauxAssistantMessage("done"), + ]); + await harness.session.prompt("Propose an approach."); + + const results = harness.session.messages.filter((message) => message.role === "toolResult"); + const blocked = results.find((message) => message.toolName === "write_file"); + expect(blocked?.isError).toBe(true); + expect(JSON.stringify(blocked?.content)).toContain("Plan mode is active"); + // Nothing was removed: the guard, not the tool list, enforces plan mode. + expect(harness.session.agent.state.tools.map((tool) => tool.name)).toEqual(toolsBefore); + } finally { + harness.cleanup(); + } +});