fix: make undo plan-aware

.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.

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<ContextMessage>();
     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;
   }

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<void> {
     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();
   }
 

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<void> {