From 9662e76afd8ae9a29ae5197db0fe3b09d1a9d4a5 Mon Sep 17 00:00:00 2001 From: liuedcson <332840681+liuedcson@users.noreply.github.com> Date: Fri, 23 Oct 2026 12:32:58 +0800 Subject: [PATCH 1/3] fix(coding-agent): keep the tool block stable across plan mode Plan mode narrowed the active tool list by dropping the editing aliases and re-appending the plan tools. A real Step profile already offers every plan tool, so that rewrite produced the same set in a different order: 42 tools, 21 positions changed, no capability gained or lost. The provider serializes `tools` ahead of the system prompt and messages, so the request's whole cached prefix is keyed on that block. Entering plan mode therefore invalidated the cache for the entire conversation and re-billed the full prompt. Measured on a live session, one entry at a 233k-token prompt read back 1.6k tokens (0.7%), and an agent-initiated entry cost another 42k. Plan mode is a permission narrowing, not a capability change. The `tool_call` hook already denies mutations outside the plan file, so the editing tools do not need to disappear. Leave the active list untouched when it already covers the plan tools; only a profile that genuinely lacks one gets it appended, which the exit path reverts. Exiting and branch restore also skip reapplying an identical list so a byte-identical prefix is never rebuilt. Covered by a suite regression that captures the request tool block and system prompt across an entry, and by extension tests for the append-only case, the unchanged-order case, and the remaining write guard. --- docs/plan-approval.md | 11 +- .../src/features/plan-mode-tools.ts | 2 +- .../coding-agent/src/features/step-plan.ts | 52 ++++++-- .../test/step-plan-extension.test.ts | 106 +++++++++++++++- .../plan-mode-prompt-cache-prefix.test.ts | 115 ++++++++++++++++++ 5 files changed, 271 insertions(+), 15 deletions(-) create mode 100644 packages/coding-agent/test/suite/regressions/plan-mode-prompt-cache-prefix.test.ts 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/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..f97c811d 100644 --- a/packages/coding-agent/src/features/step-plan.ts +++ b/packages/coding-agent/src/features/step-plan.ts @@ -45,6 +45,10 @@ export type StepPlanSource = "user" | "agent" | "unknown"; interface PlanState { enabled: boolean; + /** + * Baseline restored when plan mode leaves. Written only when entering had to + * add a tool, plus by older builds that narrowed the active list on entry. + */ toolsBeforePlanMode?: string[]; planFilePath?: string; /** "user" for /plan or --plan, "agent" for enter_plan_mode. */ @@ -55,6 +59,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); @@ -112,14 +121,35 @@ export const createStepPlanExtension = const persistPlanState = (): void => { pi.appendEntry("step-plan", { enabled, toolsBeforePlanMode, planFilePath, planSource }); }; + /** + * 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. Removing the editing tools would + * mean rewriting 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; + toolsBeforePlanMode ??= active; + pi.setActiveTools([...active, ...missing]); }; const restoreTools = (): void => { - if (toolsBeforePlanMode) pi.setActiveTools(toolsBeforePlanMode); + if (!toolsBeforePlanMode) return; + const baseline = toolsBeforePlanMode; toolsBeforePlanMode = undefined; + // An unchanged list would rebuild the same system prompt and tool block; + // skip it so leaving plan mode cannot invalidate a byte-identical prefix. + if (!sameToolOrder(pi.getActiveTools(), baseline)) pi.setActiveTools(baseline); }; const resolvePlanFilePath = (ctx: ExtensionContext): string => planFilePath ?? getPlanFilePath(ctx.sessionManager.getSessionId(), ctx.cwd); @@ -218,7 +248,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,17 +258,21 @@ export const createStepPlanExtension = : [], ); const restoredPlanState = planStates[0]; - // Off snapshots omit transient tools; a preceding plan entry retains the baseline. + // A baseline comes from an entry that had to add a plan tool, or from a + // session written by an older build, which narrowed the active list on + // entry. That branch owns the list, so drop tools added since. const savedTools = planStates.find((state) => state?.toolsBeforePlanMode)?.toolsBeforePlanMode; enabled = restoredPlanState?.enabled ?? false; - toolsBeforePlanMode = savedTools ? unique([...savedTools, ...activePlanTaskTools]) : undefined; planFilePath = restoredPlanState?.planFilePath; planSource = enabled ? (restoredPlanState?.planSource ?? "unknown") : undefined; + if (savedTools) { + const baseline = unique([...savedTools, ...activePlanTaskTools]); + if (!sameToolOrder(pi.getActiveTools(), baseline)) pi.setActiveTools(baseline); + } + toolsBeforePlanMode = undefined; if (enabled) { planFilePath = resolvePlanFilePath(ctx); enable(); - } else { - restoreTools(); } if (restoredPlanState) { migrateLegacyPlanTodos({ diff --git a/packages/coding-agent/test/step-plan-extension.test.ts b/packages/coding-agent/test/step-plan-extension.test.ts index 73ef5ee3..656f5b18 100644 --- a/packages/coding-agent/test/step-plan-extension.test.ts +++ b/packages/coding-agent/test/step-plan-extension.test.ts @@ -664,8 +664,12 @@ test.each(planEntryCases)( 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"); @@ -706,6 +710,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" 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..ff1f9602 --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/plan-mode-prompt-cache-prefix.test.ts @@ -0,0 +1,115 @@ +import { fauxAssistantMessage, fauxToolCall } from "@step-harness/providers"; +import { Type } from "typebox"; +import { expect, it } from "vitest"; +import { createStepPlanExtension } from "../../../src/features/step-plan.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 ~230k prompt tokens read back only 1.6k. + */ + +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) { + return { + name, + label: name, + description: `${name} tool`, + parameters: Type.Object({}), + execute: async () => ({ content: [{ type: "text" as const, text: name }], details: {} }), + }; +} + +interface CapturedRequest { + tools: string[]; + systemPrompt: string; +} + +it("keeps the request tool block and system prompt byte-identical across a plan-mode entry", async () => { + const harness = await createHarness({ + tools: [...PLAN_TOOLS, "find_tools"].map(stepTool), + extensionFactories: [{ name: "step-plan", factory: createStepPlanExtension(), hidden: true }], + }); + + try { + // This is the shape of an ordinary Step profile: every plan tool is already + // available, so plan mode has nothing to add. + harness.session.setActiveToolsByName([...PLAN_TOOLS, "find_tools", "enter_plan_mode", "exit_plan_mode"]); + + const captured: CapturedRequest[] = []; + harness.setResponses([ + (context) => { + captured.push({ + tools: (context.tools ?? []).map((tool) => tool.name), + systemPrompt: context.systemPrompt ?? "", + }); + return fauxAssistantMessage(fauxToolCall("enter_plan_mode", {}), { stopReason: "toolUse" }); + }, + (context) => { + captured.push({ + tools: (context.tools ?? []).map((tool) => tool.name), + systemPrompt: context.systemPrompt ?? "", + }); + return fauxAssistantMessage("done"); + }, + ]); + + await harness.session.prompt("Propose an approach."); + + expect(captured).toHaveLength(2); + // Same tools, same order, same prompt: the cached prefix survives. + expect(captured[1]?.tools).toEqual(captured[0]?.tools); + expect(captured[1]?.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(stepTool), + 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).toBeDefined(); + expect(blocked?.isError).toBe(true); + expect(JSON.stringify(blocked?.content)).toContain("Only the plan file may be written"); + // 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(); + } +}); From 1fe3655d97e758f12ab8dd8320ed4a6660d5484d Mon Sep 17 00:00:00 2001 From: liuedcson <332840681+liuedcson@users.noreply.github.com> Date: Fri, 23 Oct 2026 12:45:42 +0800 Subject: [PATCH 2/3] fix(coding-agent): scope plan-mode tool restore to its own entry Review follow-ups for the cache-stable tool block. An entry that added nothing stopped writing `toolsBeforePlanMode`, but restore still searched every entry for one. A later plan session could therefore replay an older session's snapshot: the list was reordered and tools gained since were dropped, which is the same prefix invalidation this change set out to remove. Entries now carry `version: 2`, and a baseline is only replayed when the entry that narrowed the list is a legacy one. Exit now subtracts the tools entering actually appended instead of replaying a snapshot, so a tool activated mid-plan (a late MCP catalog, find_tools) keeps its place. The plan/session guard resolves both paths through `resolveToCwd`, the same helper the file tools use, so `~`, `@`, and unicode-space inputs cannot pass the check and then land elsewhere. The denial now states that plan mode is a permission restriction, forbids retrying or routing around it through a shell, and points at the plan file and `exit_plan_mode`. Tests: byte-identical prefix across a full entry/write/exit cycle (previously only entry was covered), no persisted baseline when entry adds nothing, no reapply on a new-format restore, a stale legacy baseline not replayed, and a mid-plan tool surviving exit. All 17 plan-mode cases fail against `main` and pass here. --- .../coding-agent/docs/step-integration.md | 7 +- .../coding-agent/src/features/step-plan.ts | 70 ++++++---- .../test/step-plan-extension.test.ts | 125 +++++++++++++++++- .../plan-mode-prompt-cache-prefix.test.ts | 83 ++++++++---- 4 files changed, 230 insertions(+), 55 deletions(-) diff --git a/packages/coding-agent/docs/step-integration.md b/packages/coding-agent/docs/step-integration.md index 527b8702..f2bae90e 100644 --- a/packages/coding-agent/docs/step-integration.md +++ b/packages/coding-agent/docs/step-integration.md @@ -415,8 +415,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/step-plan.ts b/packages/coding-agent/src/features/step-plan.ts index f97c811d..9852de0a 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,10 +44,14 @@ 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; /** - * Baseline restored when plan mode leaves. Written only when entering had to - * add a tool, plus by older builds that narrowed the active list on entry. + * 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; @@ -111,7 +115,8 @@ 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 @@ -119,17 +124,19 @@ export const createStepPlanExtension = const events = (pi as { events?: EventBus }).events; const persistPlanState = (): void => { - pi.appendEntry("step-plan", { enabled, toolsBeforePlanMode, planFilePath, planSource }); + // No tool snapshot: this build adds tools rather than replacing the list, + // so `version` also tells a later restore that the branch never narrowed it. + pi.appendEntry("step-plan", { version: 2, enabled, planFilePath, planSource }); }; /** * 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. Removing the editing tools would - * mean rewriting 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. + * 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 @@ -140,16 +147,18 @@ export const createStepPlanExtension = const active = pi.getActiveTools(); const missing = PLAN_TOOLS.filter((name) => !active.includes(name)); if (missing.length === 0) return; - toolsBeforePlanMode ??= active; + planModeAddedTools = missing; pi.setActiveTools([...active, ...missing]); }; const restoreTools = (): void => { - if (!toolsBeforePlanMode) return; - const baseline = toolsBeforePlanMode; - toolsBeforePlanMode = undefined; - // An unchanged list would rebuild the same system prompt and tool block; - // skip it so leaving plan mode cannot invalidate a byte-identical prefix. - if (!sameToolOrder(pi.getActiveTools(), baseline)) pi.setActiveTools(baseline); + 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); @@ -217,7 +226,9 @@ export const createStepPlanExtension = const planPath = resolvePlanFilePath(ctx); 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)); @@ -225,7 +236,12 @@ 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, rephrase the path, or work around it with a shell command.", + `Write the implementation proposal to the plan file with write_file, then call exit_plan_mode to submit it for review.`, + ].join("\n"), }; }); @@ -258,18 +274,20 @@ export const createStepPlanExtension = : [], ); const restoredPlanState = planStates[0]; - // A baseline comes from an entry that had to add a plan tool, or from a - // session written by an older build, which narrowed the active list on - // entry. That branch owns the list, so drop tools added since. - const savedTools = planStates.find((state) => state?.toolsBeforePlanMode)?.toolsBeforePlanMode; + // The list to undo belongs to whichever entry narrowed it: a legacy + // on-entry recorded what it replaced, and so did an older off-snapshot + // that kept the baseline. This build records nothing, so its own entries + // (marked by `version`) never replay an older plan session's snapshot — + // doing that would reorder and drop tools the branch never touched. + const governingEntry = planStates.find((state) => state?.enabled || state?.toolsBeforePlanMode !== undefined); + const legacyBaseline = governingEntry?.version ? undefined : governingEntry?.toolsBeforePlanMode; enabled = restoredPlanState?.enabled ?? false; planFilePath = restoredPlanState?.planFilePath; planSource = enabled ? (restoredPlanState?.planSource ?? "unknown") : undefined; - if (savedTools) { - const baseline = unique([...savedTools, ...activePlanTaskTools]); + if (legacyBaseline) { + const baseline = unique([...legacyBaseline, ...activePlanTaskTools]); if (!sameToolOrder(pi.getActiveTools(), baseline)) pi.setActiveTools(baseline); } - toolsBeforePlanMode = undefined; if (enabled) { planFilePath = resolvePlanFilePath(ctx); enable(); diff --git a/packages/coding-agent/test/step-plan-extension.test.ts b/packages/coding-agent/test/step-plan-extension.test.ts index 656f5b18..0560f4ec 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,11 +668,13 @@ 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"]), ); @@ -919,14 +933,119 @@ 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); + 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 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) => { @@ -954,9 +1073,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 index ff1f9602..51016f13 100644 --- 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 @@ -1,7 +1,11 @@ +import { mkdirSync, writeFileSync } from "node:fs"; +import path from "node:path"; import { fauxAssistantMessage, fauxToolCall } from "@step-harness/providers"; import { Type } from "typebox"; -import { expect, it } from "vitest"; +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"; /** @@ -11,7 +15,7 @@ import { createHarness } from "../harness.ts"; * * 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 ~230k prompt tokens read back only 1.6k. + * one entry at a ~230k-token prompt read back only 1.6k tokens. */ const PLAN_TOOLS = [ @@ -25,56 +29,86 @@ const PLAN_TOOLS = [ "edit_file", ] as const; -function stepTool(name: string) { +function stepTool(name: string, execute?: (args: Record) => void) { return { name, label: name, description: `${name} tool`, - parameters: Type.Object({}), - execute: async () => ({ content: [{ type: "text" as const, text: name }], details: {} }), + 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 entry", async () => { +it("keeps the request tool block and system prompt byte-identical across a plan-mode cycle", async () => { + const cwd = process.cwd(); const harness = await createHarness({ - tools: [...PLAN_TOOLS, "find_tools"].map(stepTool), + tools: [...PLAN_TOOLS, "find_tools"].map((name) => + stepTool(name, (args) => { + if (name !== "write_file" || typeof args.path !== "string") return; + const target = path.resolve(cwd, args.path); + mkdirSync(path.dirname(target), { recursive: true }); + writeFileSync(target, typeof args.content === "string" ? args.content : ""); + }), + ), extensionFactories: [{ name: "step-plan", factory: createStepPlanExtension(), hidden: true }], }); try { - // This is the shape of an ordinary Step profile: every plan tool is already - // available, so plan mode has nothing to add. + // 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) => { - captured.push({ - tools: (context.tools ?? []).map((tool) => tool.name), - systemPrompt: context.systemPrompt ?? "", - }); + capture(context); return fauxAssistantMessage(fauxToolCall("enter_plan_mode", {}), { stopReason: "toolUse" }); }, (context) => { - captured.push({ - tools: (context.tools ?? []).map((tool) => tool.name), - systemPrompt: context.systemPrompt ?? "", - }); + capture(context); + return fauxAssistantMessage( + fauxToolCall("write_file", { path: path.relative(cwd, 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."); - expect(captured).toHaveLength(2); - // Same tools, same order, same prompt: the cached prefix survives. - expect(captured[1]?.tools).toEqual(captured[0]?.tools); - expect(captured[1]?.systemPrompt).toBe(captured[0]?.systemPrompt); + // The system prompt embeds the active tool snippets, so equal tools plus an + // equal prompt is exactly the prefix the provider caches. + console.log("DBG", JSON.stringify(captured.map((c) => c.tools))); + 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(); @@ -82,8 +116,9 @@ it("keeps the request tool block and system prompt byte-identical across a plan- }); it("still blocks non-plan mutations after entering plan mode with tools intact", async () => { + const cwd = process.cwd(); const harness = await createHarness({ - tools: [...PLAN_TOOLS, "find_tools"].map(stepTool), + tools: [...PLAN_TOOLS, "find_tools"].map((name) => stepTool(name)), extensionFactories: [{ name: "step-plan", factory: createStepPlanExtension(), hidden: true }], }); @@ -104,11 +139,11 @@ it("still blocks non-plan mutations after entering plan mode with tools intact", const results = harness.session.messages.filter((message) => message.role === "toolResult"); const blocked = results.find((message) => message.toolName === "write_file"); - expect(blocked).toBeDefined(); expect(blocked?.isError).toBe(true); - expect(JSON.stringify(blocked?.content)).toContain("Only the plan file may be written"); + 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); + void cwd; } finally { harness.cleanup(); } From 6b810d1d18470e125e0273a047eacac4a8b0bf8e Mon Sep 17 00:00:00 2001 From: liuedcson <332840681+liuedcson@users.noreply.github.com> Date: Fri, 23 Oct 2026 13:15:41 +0800 Subject: [PATCH 3/3] fix(coding-agent): harden the plan-mode guard and legacy restore MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups. The plan-mode tool block is now the only thing that changes nothing while the guard does all the enforcement, so the guard needs coverage and the legacy restore needs to stop losing state. - Cover the path normalization the guard depends on. Reverting the comparison to `path.resolve` previously kept every test green; the entry loop now accepts `@`-prefixed, absolute, and `..`-spelled forms of the plan path and rejects a leading-space path, an `@`-prefixed non-plan path, and a traversal out of the plan file for all four editing tools. - Carry a legacy tool snapshot through a todo migration. Rewriting the entry used to drop the only record of the list plan mode narrowed, which stranded the branch on the next restore. - Replay only the baseline the governing entry itself carries, so a newer plan session never inherits an earlier one's list. - Align the denial with the documented policy: `run_command` keeps its normal permissions, so the message tells the model to stay read-only instead of claiming a shell workaround is blocked. - Guard a null tool input, and stop the docs from claiming every file-writing tool is covered — the check names the four editing tools, and MCP tools are outside it. - Remove a leftover debug log and base the regression's working directory on the harness workspace rather than the test process. All 184 plan-related cases pass; `tsgo --noEmit`, `biome check`, `check:ts-imports`, `check:layer-direction`, `check:public-boundary`, and `check:coding-agent-entry-freeze` are clean. --- .../coding-agent/docs/step-integration.md | 12 ++-- .../coding-agent/src/features/step-plan.ts | 40 ++++++++---- .../test/step-plan-extension.test.ts | 65 +++++++++++++++++-- .../plan-mode-prompt-cache-prefix.test.ts | 14 ++-- 4 files changed, 104 insertions(+), 27 deletions(-) diff --git a/packages/coding-agent/docs/step-integration.md b/packages/coding-agent/docs/step-integration.md index f2bae90e..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 diff --git a/packages/coding-agent/src/features/step-plan.ts b/packages/coding-agent/src/features/step-plan.ts index 9852de0a..62e01742 100644 --- a/packages/coding-agent/src/features/step-plan.ts +++ b/packages/coding-agent/src/features/step-plan.ts @@ -123,10 +123,23 @@ export const createStepPlanExtension = // ExtensionAPI keep working without one. const events = (pi as { events?: EventBus }).events; - const persistPlanState = (): void => { - // No tool snapshot: this build adds tools rather than replacing the list, - // so `version` also tells a later restore that the branch never narrowed it. - pi.appendEntry("step-plan", { version: 2, enabled, 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. @@ -224,7 +237,7 @@ 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 : ""; // Resolve exactly as the file tools do (`~`, `@`, unicode spaces), so an // allowed target cannot land somewhere else than the guard checked. @@ -239,7 +252,8 @@ export const createStepPlanExtension = 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, rephrase the path, or work around it with a shell command.", + "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"), }; @@ -274,13 +288,13 @@ export const createStepPlanExtension = : [], ); const restoredPlanState = planStates[0]; - // The list to undo belongs to whichever entry narrowed it: a legacy - // on-entry recorded what it replaced, and so did an older off-snapshot - // that kept the baseline. This build records nothing, so its own entries - // (marked by `version`) never replay an older plan session's snapshot — - // doing that would reorder and drop tools the branch never touched. + // 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?.version ? undefined : governingEntry?.toolsBeforePlanMode; + const legacyBaseline = governingEntry?.toolsBeforePlanMode; enabled = restoredPlanState?.enabled ?? false; planFilePath = restoredPlanState?.planFilePath; planSource = enabled ? (restoredPlanState?.planSource ?? "unknown") : undefined; @@ -297,7 +311,7 @@ export const createStepPlanExtension = 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 0560f4ec..d584f188 100644 --- a/packages/coding-agent/test/step-plan-extension.test.ts +++ b/packages/coding-agent/test/step-plan-extension.test.ts @@ -688,11 +688,32 @@ test.each(planEntryCases)( 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 }); @@ -949,6 +970,11 @@ test("exiting plan mode removes only the tools entering added", async () => { // 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); @@ -964,6 +990,37 @@ test("exiting plan mode removes only the tools entering added", async () => { 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"]; 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 index 51016f13..b1d48425 100644 --- 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 @@ -51,18 +51,20 @@ interface CapturedRequest { } it("keeps the request tool block and system prompt byte-identical across a plan-mode cycle", async () => { - const cwd = process.cwd(); + // 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(cwd, args.path); + 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, @@ -85,7 +87,10 @@ it("keeps the request tool block and system prompt byte-identical across a plan- (context) => { capture(context); return fauxAssistantMessage( - fauxToolCall("write_file", { path: path.relative(cwd, planPath), content: "# Plan\n\n1. Do it.\n" }), + fauxToolCall("write_file", { + path: path.relative(workspace, planPath), + content: "# Plan\n\n1. Do it.\n", + }), { stopReason: "toolUse" }, ); }, @@ -103,7 +108,6 @@ it("keeps the request tool block and system prompt byte-identical across a plan- // The system prompt embeds the active tool snippets, so equal tools plus an // equal prompt is exactly the prefix the provider caches. - console.log("DBG", JSON.stringify(captured.map((c) => c.tools))); expect(captured).toHaveLength(4); for (const request of captured.slice(1)) { expect(request.tools).toEqual(captured[0]?.tools); @@ -116,7 +120,6 @@ it("keeps the request tool block and system prompt byte-identical across a plan- }); it("still blocks non-plan mutations after entering plan mode with tools intact", async () => { - const cwd = process.cwd(); const harness = await createHarness({ tools: [...PLAN_TOOLS, "find_tools"].map((name) => stepTool(name)), extensionFactories: [{ name: "step-plan", factory: createStepPlanExtension(), hidden: true }], @@ -143,7 +146,6 @@ it("still blocks non-plan mutations after entering plan mode with tools intact", 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); - void cwd; } finally { harness.cleanup(); }