From 4ccfc26594d5d4906acfcf24704dde2091403029 Mon Sep 17 00:00:00 2001 From: qer Date: Thu, 2 Jul 2026 21:53:32 +0800 Subject: [PATCH] fix: make undo plan-aware --- .changeset/undo-plan-mode.md | 5 ++ .../agent-core/src/agent/context/index.ts | 8 ++ packages/agent-core/src/agent/plan/index.ts | 30 +++++++ packages/agent-core/test/agent/plan.test.ts | 86 +++++++++++++++++++ 4 files changed, 129 insertions(+) create mode 100644 .changeset/undo-plan-mode.md diff --git a/.changeset/undo-plan-mode.md b/.changeset/undo-plan-mode.md new file mode 100644 index 000000000..158848b65 --- /dev/null +++ b/.changeset/undo-plan-mode.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Make /undo exit plan mode when it undoes the turn that entered plan mode. diff --git a/packages/agent-core/src/agent/context/index.ts b/packages/agent-core/src/agent/context/index.ts index b9192ded0..aeb0a341b 100644 --- a/packages/agent-core/src/agent/context/index.ts +++ b/packages/agent-core/src/agent/context/index.ts @@ -157,6 +157,7 @@ export class ContextMemory { this._lastAssistantAt = null; this.agent.microCompaction.reset(); this.agent.injection.onContextClear(); + this.agent.planMode.onContextClear(); this.agent.emitStatusUpdated(); } @@ -167,6 +168,7 @@ export class ContextMemory { this.agent.records.logRecord({ type: 'context.undo', count }); let removedUserCount = 0; + let oldestRemovedUserIndex: number | undefined; const removedMessages = new Set(); let stoppedAtBoundary = false; for (let i = this._history.length - 1; i >= 0; i--) { @@ -189,6 +191,7 @@ export class ContextMemory { if (isRealUserInput(message)) { removedUserCount++; + oldestRemovedUserIndex = i; if (removedUserCount >= count) break; } } @@ -218,6 +221,10 @@ export class ContextMemory { }, ); } + + if (oldestRemovedUserIndex !== undefined) { + this.agent.planMode.cancelAfterUndoIfNeeded(oldestRemovedUserIndex); + } } applyCompaction(input: CompactionInput): CompactionResult { @@ -297,6 +304,7 @@ export class ContextMemory { this.tokenCountCoveredMessageCount = this._history.length; this.agent.microCompaction.reset(); this.agent.injection.onContextCompacted(); + this.agent.planMode.onContextCompacted(); this.agent.emitStatusUpdated(); return result; } diff --git a/packages/agent-core/src/agent/plan/index.ts b/packages/agent-core/src/agent/plan/index.ts index fdaafdbb1..d9eba05a6 100644 --- a/packages/agent-core/src/agent/plan/index.ts +++ b/packages/agent-core/src/agent/plan/index.ts @@ -15,6 +15,7 @@ export class PlanMode { protected _isActive = false; protected _planId: null | string = null; protected _planFilePath: PlanFilePath = null; + protected _enteredAtHistoryLength: number | null = null; constructor(protected readonly agent: Agent) {} @@ -30,6 +31,7 @@ export class PlanMode { this._isActive = true; this._planId = id; this._planFilePath = null; + this._enteredAtHistoryLength = null; let enterRecorded = false; try { @@ -41,6 +43,7 @@ export class PlanMode { if (createFile) { await this.writeEmptyPlanFile(planFilePath); } + this._enteredAtHistoryLength = this.agent.context.history.length; } catch (error) { if (enterRecorded) { this.cancel(id); @@ -48,6 +51,7 @@ export class PlanMode { this._isActive = false; this._planId = null; this._planFilePath = null; + this._enteredAtHistoryLength = null; } throw error; } @@ -64,6 +68,7 @@ export class PlanMode { this._isActive = true; this._planId = id; this._planFilePath = this.planFilePathFor(id); + this._enteredAtHistoryLength = this.agent.context.history.length; } cancel(id?: string): void { @@ -75,9 +80,33 @@ export class PlanMode { this._isActive = false; this._planId = null; this._planFilePath = null; + this._enteredAtHistoryLength = null; this.agent.emitStatusUpdated(); } + cancelAfterUndoIfNeeded(removedUserIndex: number): void { + if (!this._isActive) return; + if (this.agent.records.restoring !== null) return; + if (this._enteredAtHistoryLength === null) return; + // Cancel only when plan mode was entered after the user prompt being + // undone. Manual `/plan on` before a prompt should survive undoing that prompt. + if (this._enteredAtHistoryLength > removedUserIndex) { + this.cancel(this._planId ?? undefined); + } + } + + onContextClear(): void { + if (this._isActive && this._enteredAtHistoryLength !== null) { + this._enteredAtHistoryLength = 0; + } + } + + onContextCompacted(): void { + if (this._isActive && this._enteredAtHistoryLength !== null) { + this._enteredAtHistoryLength = 0; + } + } + async clear(): Promise { if (!this._planFilePath) return; await this.writeEmptyPlanFile(this._planFilePath); @@ -92,6 +121,7 @@ export class PlanMode { this._isActive = false; this._planId = null; this._planFilePath = null; + this._enteredAtHistoryLength = null; this.agent.emitStatusUpdated(); } diff --git a/packages/agent-core/test/agent/plan.test.ts b/packages/agent-core/test/agent/plan.test.ts index 9996e7829..abd6bb991 100644 --- a/packages/agent-core/test/agent/plan.test.ts +++ b/packages/agent-core/test/agent/plan.test.ts @@ -597,6 +597,92 @@ describe('plan mode injection cadence', () => { expect(lastUserText(ctx.agent.context.history)).toContain('Plan mode is active'); }); + + it('cancels plan mode when undo removes the turn that entered planning', async () => { + const enterPlanModeCall: ToolCall = { + type: 'function', + id: 'call_enter_plan_undo', + name: 'EnterPlanMode', + arguments: '{}', + }; + const ctx = testAgent({ kaos: createPlanKaos() }); + ctx.configure({ tools: ['EnterPlanMode'] }); + await ctx.rpc.setPermission({ mode: 'yolo' }); + + ctx.mockNextResponse({ type: 'text', text: 'I will enter plan mode.' }, enterPlanModeCall); + ctx.mockNextResponse({ type: 'text', text: 'Plan mode is active now.' }); + await ctx.rpc.prompt({ input: [{ type: 'text', text: 'Plan first' }] }); + await ctx.untilTurnEnd(); + + ctx.newEvents(); + await ctx.rpc.undoHistory({ count: 1 }); + + expect(ctx.agent.planMode.isActive).toBe(false); + await expect(ctx.rpc.getPlan({})).resolves.toBeNull(); + expect(ctx.newEvents()).toContainEqual( + expect.objectContaining({ + type: '[wire]', + event: 'context.undo', + args: expect.objectContaining({ count: 1 }), + }), + ); + expect(ctx.newEvents()).toContainEqual( + expect.objectContaining({ + type: '[wire]', + event: 'plan_mode.cancel', + }), + ); + }); + + it('keeps plan mode active when undo stops after the plan-enter anchor', async () => { + const ctx = testAgent({ kaos: createPlanKaos() }); + ctx.configure(); + await ctx.rpc.enterPlan({}); + + ctx.agent.context.appendUserMessage([{ type: 'text', text: 'first plan request' }]); + await ctx.agent.injection.inject(); + ctx.appendAssistantTurn(1, 'First plan drafted.'); + + ctx.agent.context.appendUserMessage([{ type: 'text', text: 'second plan request' }]); + await ctx.agent.injection.inject(); + ctx.appendAssistantTurn(2, 'Second plan drafted.'); + + await ctx.rpc.undoHistory({ count: 1 }); + + expect(ctx.agent.planMode.isActive).toBe(true); + await expect(ctx.rpc.getPlan({})).resolves.toMatchObject({ + path: ctx.agent.planMode.planFilePath, + }); + }); + + it('restores plan mode off from an undo followed by plan cancel', () => { + const ctx = testAgent(); + ctx.configure(); + + ctx.dispatch({ + type: 'plan_mode.enter', + id: 'restored-plan', + }); + ctx.dispatch({ + type: 'context.append_message', + message: { + role: 'user', + content: [{ type: 'text', text: 'draft the plan' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + }); + ctx.dispatch({ + type: 'context.undo', + count: 1, + }); + ctx.dispatch({ + type: 'plan_mode.cancel', + id: 'restored-plan', + }); + + expect(ctx.agent.planMode.isActive).toBe(false); + }); }); function delay(ms: number): Promise {