diff --git a/.changeset/undo-rolls-back-todos.md b/.changeset/undo-rolls-back-todos.md new file mode 100644 index 0000000000..1cbf0bcb9e --- /dev/null +++ b/.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. diff --git a/packages/agent-core/src/agent/context/index.ts b/packages/agent-core/src/agent/context/index.ts index 52f8c6c4fb..8d06ab8c94 100644 --- a/packages/agent-core/src/agent/context/index.ts +++ b/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(); + } } } } diff --git a/packages/agent-core/src/agent/index.ts b/packages/agent-core/src/agent/index.ts index c54bd0896b..f3087ecd13 100644 --- a/packages/agent-core/src/agent/index.ts +++ b/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), }; } diff --git a/packages/agent-core/src/agent/tool/index.ts b/packages/agent-core/src/agent/tool/index.ts index e5874018e8..6cf78fbe28 100644 --- a/packages/agent-core/src/agent/tool/index.ts +++ b/packages/agent-core/src/agent/tool/index.ts @@ -72,6 +72,19 @@ export class ToolManager { */ private readonly pendingLoadedDynamicTools = new Set(); protected readonly store: Partial = {}; + /** + * 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> = [{}]; 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(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(); } /** diff --git a/packages/agent-core/src/rpc/core-api.ts b/packages/agent-core/src/rpc/core-api.ts index 1ae6239c7a..df9013fa58 100644 --- a/packages/agent-core/src/rpc/core-api.ts +++ b/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[]; } diff --git a/packages/agent-core/src/rpc/core-impl.ts b/packages/agent-core/src/rpc/core-impl.ts index 17b261aa30..1cbf385357 100644 --- a/packages/agent-core/src/rpc/core-impl.ts +++ b/packages/agent-core/src/rpc/core-impl.ts @@ -1608,6 +1608,10 @@ export class KimiCore implements PromisableMethods { return this.sessionApi(sessionId).getTools(payload); } + getTodos({ sessionId, ...payload }: SessionAgentPayload) { + return this.sessionApi(sessionId).getTodos(payload); + } + getBackground({ sessionId, ...payload }: SessionAgentPayload) { return this.sessionApi(sessionId).getBackground(payload); } diff --git a/packages/agent-core/src/session/rpc.ts b/packages/agent-core/src/session/rpc.ts index d6ee10d236..6ee98edd53 100644 --- a/packages/agent-core/src/session/rpc.ts +++ b/packages/agent-core/src/session/rpc.ts @@ -317,6 +317,10 @@ export class SessionAPIImpl implements PromisableMethods { return (await this.getAgent(agentId)).getTools(payload); } + async getTodos({ agentId, ...payload }: AgentScopedPayload) { + return (await this.getAgent(agentId)).getTodos(payload); + } + async getBackground({ agentId, ...payload }: AgentScopedPayload) { return (await this.getAgent(agentId)).getBackground(payload); } diff --git a/packages/agent-core/test/agent/context.test.ts b/packages/agent-core/test/agent/context.test.ts index 432187410a..da6f33f2ae 100644 --- a/packages/agent-core/test/agent/context.test.ts +++ b/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', () => { diff --git a/packages/agent-core/test/agent/resume.test.ts b/packages/agent-core/test/agent/resume.test.ts index 0bb76debe9..4e915364e5 100644 --- a/packages/agent-core/test/agent/resume.test.ts +++ b/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([ { diff --git a/packages/node-sdk/src/rpc.ts b/packages/node-sdk/src/rpc.ts index 44b35136cf..da55bac008 100644 --- a/packages/node-sdk/src/rpc.ts +++ b/packages/node-sdk/src/rpc.ts @@ -720,11 +720,11 @@ export abstract class SDKRpcClientBase { } async getTodos(input: SessionIdRpcInput): Promise { - void input; - throw new KimiError( - ErrorCodes.NOT_IMPLEMENTED, - 'getTodos is only available on the agent-core-v2 engine.', - ); + const rpc = await this.getRpc(); + return rpc.getTodos({ + sessionId: input.sessionId, + agentId: this.interactiveAgentId, + }); } async undoHistory(input: SessionIdRpcInput & { count: number }): Promise { diff --git a/packages/node-sdk/test/session-context.test.ts b/packages/node-sdk/test/session-context.test.ts index 36f9564d09..85d0db3000 100644 --- a/packages/node-sdk/test/session-context.test.ts +++ b/packages/node-sdk/test/session-context.test.ts @@ -67,6 +67,28 @@ describe('Session context', () => { } }); + it('serves getTodos on the v1 engine and keeps it callable after undo', async () => { + const homeDir = await makeTempDir(tempDirs, 'kimi-sdk-context-todos-home-'); + const workDir = await makeTempDir(tempDirs, 'kimi-sdk-context-todos-work-'); + await writeTestConfig(homeDir, 200_000); + const harness = createKimiHarness({ homeDir, identity: TEST_IDENTITY }); + + try { + const session = await harness.createSession({ id: 'ses_context_todos', workDir }); + await session.importContext('Earlier context to undo.', "file 'notes.md'"); + + await expect(session.getTodos()).resolves.toEqual([]); + + await session.undoHistory(1); + + // The TUI's post-undo todo panel refresh relies on this resolving + // instead of throwing NOT_IMPLEMENTED. + await expect(session.getTodos()).resolves.toEqual([]); + } finally { + await harness.close(); + } + }); + it('appends old-compatible user context markup when importing raw content', async () => { const homeDir = await makeTempDir(tempDirs, 'kimi-sdk-context-import-home-'); const workDir = await makeTempDir(tempDirs, 'kimi-sdk-context-import-work-');