fix(agent-core): roll the todo store back on undo

.changeset/undo-rolls-back-todos.md

@@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": patch
+---
+
+The todo list now rolls back when you undo prompts, and stays rolled back after resuming the session.

packages/agent-core/src/agent/context/index.ts

@@ -285,6 +285,10 @@ export class ContextMemory {
     }
 
     this.agent.replayBuilder.removeLastMessages(removedMessages);
+    // Roll the tool store back to match the spliced history before any
+    // undo_limit throw — a partial undo must leave the store consistent with
+    // the partially removed history, compensating records included.
+    this.agent.tools.rollbackStore(removedUserCount);
 
     this.openSteps.clear();
     this.pendingToolResultIds.clear();
@@ -798,6 +802,14 @@ export class ContextMemory {
         type: 'message',
         message,
       });
+      if (isRealUserInput(message)) {
+        // Undo-anchor checkpoint: snapshot the tool store the way this turn
+        // found it, so `undo` can restore it alongside the history tail.
+        // `pushHistory` is the one funnel every message (live, replay, and
+        // deferred-behind-an-open-exchange) passes through, which keeps
+        // checkpoint pushes 1:1 with undo's anchor counting over `_history`.
+        this.agent.tools.snapshotToolStore();
+      }
     }
   }
 }

packages/agent-core/src/agent/index.ts

@@ -718,6 +718,7 @@ export class Agent {
       getPlan: () => this.planMode.data(),
       getUsage: () => this.usage.data(),
       getTools: () => this.tools.data(),
+      getTodos: () => this.tools.storeValue('todo') ?? [],
       getBackground: (payload) => this.background.list(payload.activeOnly ?? false, payload.limit),
     };
   }

packages/agent-core/src/agent/tool/index.ts

@@ -72,6 +72,19 @@ export class ToolManager {
    */
   private readonly pendingLoadedDynamicTools = new Set<string>();
   protected readonly store: Partial<ToolStoreData> = {};
+  /**
+   * Store snapshots taken when each real user input (undo anchor) enters
+   * history, oldest first. Index 0 is the baseline seeded at construction,
+   * compaction, and clear; `rollbackStore` pops from the top on undo.
+   * Shallow copies suffice: writers replace values wholesale (TodoList swaps
+   * the whole array), never mutate in place, so snapshots share the values.
+   * The stack length is exactly 1 + the live anchor count in history —
+   * pushes are 1:1 with anchor appends, undo pops what it removed, and
+   * clear/compaction reseed — so it is bounded by history itself and undo
+   * can never request a depth the stack does not cover. Do not add trimming:
+   * a capped window would restore the store of a turn the undo just removed.
+   */
+  private storeCheckpoints: Array<Partial<ToolStoreData>> = [{}];
   private mcpToolStatusUnsubscribe: (() => void) | undefined;
   /**
    * `serverName\nhash` keys of `mcp.tools_discovered` records already durable
@@ -133,6 +146,48 @@ export class ToolManager {
     this.store[key] = value;
   }
 
+  /** Snapshot the store for a just-appended undo anchor (real user input). */
+  snapshotToolStore(): void {
+    this.storeCheckpoints.push({ ...this.store });
+  }
+
+  /**
+   * Roll the store back over `removedUserTurns` undo anchors (v2 parity: pop
+   * that many checkpoints, restore to the snapshot the earliest removed anchor
+   * found, never below the baseline). Restores go through `updateStore`, so
+   * the live path logs compensating `tools.update_store` records — the
+   * append-only wire, the transcript reducer, and the v1-fold-over-v2-wire
+   * resume all stay self-consistent — while replay restore is record-suppressed.
+   */
+  rollbackStore(removedUserTurns: number): void {
+    const pops = Math.min(removedUserTurns, this.storeCheckpoints.length - 1);
+    if (pops <= 0) return;
+    const target = this.storeCheckpoints[this.storeCheckpoints.length - pops]!;
+    this.storeCheckpoints.splice(this.storeCheckpoints.length - pops);
+    // Union of the target and current keys: a key first written during the
+    // undone turns is absent from the target and must be restored to the
+    // canonical empty (TodoList's clear form — consumers treat `undefined`
+    // and `[]` alike).
+    const keys = new Set([
+      ...(Object.keys(target) as ToolStoreKey[]),
+      ...(Object.keys(this.store) as ToolStoreKey[]),
+    ]);
+    for (const key of keys) {
+      const value = target[key];
+      if (this.store[key] === value) continue;
+      this.updateStore(key, value ?? ([] as ToolStoreData[typeof key]));
+    }
+  }
+
+  /** Typed store read for RPC surfaces (`store` itself is protected). */
+  storeValue<K extends ToolStoreKey>(key: K): ToolStoreData[K] | undefined {
+    return this.store[key];
+  }
+
+  private resetStoreCheckpoints(): void {
+    this.storeCheckpoints = [{ ...this.store }];
+  }
+
   /**
    * Execute a user-initiated `!` shell command. Reuses the builtin Bash tool
    * (same kaos / cwd / BackgroundManager as the agent), recording the command
@@ -667,6 +722,7 @@ export class ToolManager {
    */
   onContextCleared(): void {
     this.pendingLoadedDynamicTools.clear();
+    this.resetStoreCheckpoints();
   }
 
   /**
@@ -678,6 +734,7 @@ export class ToolManager {
    */
   onContextCompacted(): void {
     this.pendingLoadedDynamicTools.clear();
+    this.resetStoreCheckpoints();
   }
 
   /**

packages/agent-core/src/rpc/core-api.ts

@@ -19,6 +19,7 @@ import type { KimiConfig, KimiConfigPatch, McpServerConfig } from '#/config';
 import type { ExperimentalFeatureState } from '#/flags';
 import type { ResumeSessionResult } from '#/rpc/resumed';
 import type { SessionMeta } from '#/session';
+import type { TodoItem } from '#/tools/builtin/state/todo-list';
 import type { GlobalMcpServerConfig } from '#/mcp/global-config';
 import type { McpServerConfigView } from '#/mcp/config-view';
 import type { McpRegistryPluginOrigin, McpServerSource } from '#/mcp/registry';
@@ -639,6 +640,7 @@ export interface AgentAPI {
   getPlan: (payload: EmptyPayload) => PlanData;
   getUsage: (payload: EmptyPayload) => UsageStatus;
   getTools: (payload: EmptyPayload) => readonly ToolInfo[];
+  getTodos: (payload: EmptyPayload) => readonly TodoItem[];
   getBackground: (payload: GetBackgroundPayload) => readonly BackgroundTaskInfo[];
 }
 

packages/agent-core/src/rpc/core-impl.ts

@@ -1608,6 +1608,10 @@ export class KimiCore implements PromisableMethods<CoreAPI> {
     return this.sessionApi(sessionId).getTools(payload);
   }
 
+  getTodos({ sessionId, ...payload }: SessionAgentPayload<EmptyPayload>) {
+    return this.sessionApi(sessionId).getTodos(payload);
+  }
+
   getBackground({ sessionId, ...payload }: SessionAgentPayload<GetBackgroundPayload>) {
     return this.sessionApi(sessionId).getBackground(payload);
   }

packages/agent-core/src/session/rpc.ts

@@ -317,6 +317,10 @@ export class SessionAPIImpl implements PromisableMethods<SessionAPI> {
     return (await this.getAgent(agentId)).getTools(payload);
   }
 
+  async getTodos({ agentId, ...payload }: AgentScopedPayload<EmptyPayload>) {
+    return (await this.getAgent(agentId)).getTodos(payload);
+  }
+
   async getBackground({ agentId, ...payload }: AgentScopedPayload<GetBackgroundPayload>) {
     return (await this.getAgent(agentId)).getBackground(payload);
   }

packages/agent-core/test/agent/context.test.ts

@@ -1335,6 +1335,145 @@ describe('Agent context', () => {
     ]);
   });
 
+  it('undo rolls the tool store back to the last user prompt', async () => {
+    const ctx = testAgent();
+    ctx.configure();
+    const todosA = [{ title: 'first task', status: 'done' as const }];
+    const todosB = [{ title: 'second task', status: 'in_progress' as const }];
+
+    ctx.agent.context.appendUserMessage([{ type: 'text', text: 'first prompt' }]);
+    ctx.agent.tools.updateStore('todo', todosA);
+    ctx.agent.context.appendUserMessage([{ type: 'text', text: 'second prompt' }]);
+    ctx.agent.tools.updateStore('todo', todosB);
+
+    ctx.agent.context.undo(1);
+
+    expect(ctx.agent.tools.storeData()['todo']).toEqual(todosA);
+    await ctx.expectResumeMatches();
+  });
+
+  it('undo logs a compensating tools.update_store record on the wire', () => {
+    const ctx = testAgent();
+    ctx.configure();
+    const todosA = [{ title: 'first task', status: 'done' as const }];
+    const todosB = [{ title: 'second task', status: 'in_progress' as const }];
+
+    ctx.agent.context.appendUserMessage([{ type: 'text', text: 'first prompt' }]);
+    ctx.agent.tools.updateStore('todo', todosA);
+    ctx.agent.context.appendUserMessage([{ type: 'text', text: 'second prompt' }]);
+    ctx.agent.tools.updateStore('todo', todosB);
+    ctx.newEvents();
+
+    ctx.agent.context.undo(1);
+
+    expect(ctx.newEvents()).toContainEqual(
+      expect.objectContaining({
+        type: '[wire]',
+        event: 'tools.update_store',
+        args: expect.objectContaining({ key: 'todo', value: todosA }),
+      }),
+    );
+  });
+
+  it('undo clears todos first written during the undone turn', () => {
+    const ctx = testAgent();
+    ctx.configure();
+    const todosB = [{ title: 'second task', status: 'in_progress' as const }];
+
+    ctx.agent.context.appendUserMessage([{ type: 'text', text: 'first prompt' }]);
+    ctx.agent.tools.updateStore('todo', todosB);
+
+    ctx.agent.context.undo(1);
+
+    expect(ctx.agent.tools.storeData()['todo']).toEqual([]);
+  });
+
+  it('partial undo at the compaction boundary restores the compaction baseline', () => {
+    const ctx = testAgent();
+    ctx.configure();
+    const todosA = [{ title: 'first task', status: 'done' as const }];
+    const todosB = [{ title: 'second task', status: 'in_progress' as const }];
+
+    ctx.agent.context.appendUserMessage([{ type: 'text', text: 'old user message' }]);
+    ctx.agent.tools.updateStore('todo', todosA);
+    ctx.agent.context.applyCompaction({
+      summary: 'summary of compacted context',
+      compactedCount: 1,
+      tokensBefore: 100,
+    });
+    ctx.agent.context.appendUserMessage([{ type: 'text', text: 'recent user message' }]);
+    ctx.agent.tools.updateStore('todo', todosB);
+
+    expect(() => {
+      ctx.agent.context.undo(2);
+    }).toThrow('Cannot undo 2 prompts; only 1 prompt can be undone');
+
+    expect(ctx.agent.tools.storeData()['todo']).toEqual(todosA);
+  });
+
+  it('undo skips background notifications when rolling back the store', () => {
+    const ctx = testAgent();
+    ctx.configure();
+    const todosA = [{ title: 'first task', status: 'done' as const }];
+    const todosB = [{ title: 'second task', status: 'in_progress' as const }];
+
+    ctx.agent.context.appendUserMessage([{ type: 'text', text: 'first prompt' }]);
+    ctx.agent.tools.updateStore('todo', todosA);
+    ctx.agent.context.appendUserMessage([{ type: 'text', text: 'second prompt' }]);
+    ctx.agent.context.appendMessage({
+      role: 'user',
+      content: [{ type: 'text', text: 'background task completed' }],
+      toolCalls: [],
+      origin: {
+        kind: 'background_task',
+        taskId: 'bash-001',
+        status: 'completed',
+        notificationId: 'task:bash-001:completed',
+      },
+    });
+    ctx.agent.tools.updateStore('todo', todosB);
+
+    ctx.agent.context.undo(1);
+
+    expect(ctx.agent.tools.storeData()['todo']).toEqual(todosA);
+  });
+
+  it('clear resets the undo baseline but keeps todos', () => {
+    const ctx = testAgent();
+    ctx.configure();
+    const todosA = [{ title: 'first task', status: 'done' as const }];
+    const todosB = [{ title: 'second task', status: 'in_progress' as const }];
+
+    ctx.agent.context.appendUserMessage([{ type: 'text', text: 'first prompt' }]);
+    ctx.agent.tools.updateStore('todo', todosA);
+    ctx.agent.context.clear();
+    ctx.agent.context.appendUserMessage([{ type: 'text', text: 'post-clear prompt' }]);
+    ctx.agent.tools.updateStore('todo', todosB);
+
+    ctx.agent.context.undo(1);
+
+    expect(ctx.agent.tools.storeData()['todo']).toEqual(todosA);
+  });
+
+  it('undo deeper than a hundred anchors restores the oldest anchor snapshot', () => {
+    const ctx = testAgent();
+    ctx.configure();
+    const anchorCount = 120;
+    const todosBeforeAnchor = (i: number) => [{ title: `todos before anchor ${String(i)}`, status: 'pending' as const }];
+
+    // A distinct store value per anchor so a wrong checkpoint target shows up
+    // as the wrong value, not just a stale one.
+    for (let i = 1; i <= anchorCount; i++) {
+      ctx.agent.tools.updateStore('todo', todosBeforeAnchor(i));
+      ctx.agent.context.appendUserMessage([{ type: 'text', text: `prompt ${String(i)}` }]);
+    }
+    ctx.agent.tools.updateStore('todo', [{ title: 'final write', status: 'pending' }]);
+
+    ctx.agent.context.undo(anchorCount);
+
+    expect(ctx.agent.tools.storeData()['todo']).toEqual(todosBeforeAnchor(1));
+  });
+
 });
 
 describe('Agent context notification projection', () => {

packages/agent-core/test/agent/resume.test.ts

@@ -230,6 +230,82 @@ describe('Agent resume', () => {
     await ctx.expectResumeMatches();
   });
 
+  it('replays context.undo and restores the pre-undo tool store', async () => {
+    // Legacy wire shape: the undo record carries no compensating store write,
+    // so replay itself must fold the store back to the pre-undo value.
+    const todosA = [{ title: 'first task', status: 'done' as const }];
+    const todosB = [{ title: 'second task', status: 'in_progress' as const }];
+    const persistence = new RecordingAgentPersistence([
+      {
+        type: 'tools.update_store',
+        key: 'todo',
+        value: todosA,
+      },
+      {
+        type: 'context.append_message',
+        message: {
+          role: 'user',
+          content: [{ type: 'text', text: 'undone prompt' }],
+          toolCalls: [],
+          origin: { kind: 'user' },
+        },
+      },
+      {
+        type: 'tools.update_store',
+        key: 'todo',
+        value: todosB,
+      },
+      { type: 'context.undo', count: 1 },
+    ]);
+    const ctx = testAgent({ persistence });
+
+    await ctx.agent.resume();
+
+    expect(ctx.agent.tools.storeData()).toEqual({ todo: todosA });
+    expect(ctx.agent.context.history).toEqual([]);
+    await ctx.expectResumeMatches();
+  });
+
+  it('replays compensating store records after undo idempotently', async () => {
+    // New wire shape: the live undo already logged the compensating write.
+    const todosA = [{ title: 'first task', status: 'done' as const }];
+    const todosB = [{ title: 'second task', status: 'in_progress' as const }];
+    const persistence = new RecordingAgentPersistence([
+      {
+        type: 'tools.update_store',
+        key: 'todo',
+        value: todosA,
+      },
+      {
+        type: 'context.append_message',
+        message: {
+          role: 'user',
+          content: [{ type: 'text', text: 'undone prompt' }],
+          toolCalls: [],
+          origin: { kind: 'user' },
+        },
+      },
+      {
+        type: 'tools.update_store',
+        key: 'todo',
+        value: todosB,
+      },
+      { type: 'context.undo', count: 1 },
+      {
+        type: 'tools.update_store',
+        key: 'todo',
+        value: todosA,
+      },
+    ]);
+    const ctx = testAgent({ persistence });
+
+    await ctx.agent.resume();
+
+    expect(ctx.agent.tools.storeData()).toEqual({ todo: todosA });
+    expect(ctx.agent.context.history).toEqual([]);
+    await ctx.expectResumeMatches();
+  });
+
   it('applies wire migrations while replaying persisted records', async () => {
     const persistence = new RecordingAgentPersistence([
       {