From 8d66aa71110ed903fbcbdb6dc6eaaa607db072b3 Mon Sep 17 00:00:00 2001 From: _Kerman Date: Mon, 1 Jun 2026 18:08:39 +0800 Subject: [PATCH 01/10] feat: add /undo slash command to withdraw last prompt from history --- apps/kimi-code/src/tui/commands/dispatch.ts | 6 +++ apps/kimi-code/src/tui/commands/registry.ts | 7 ++++ apps/kimi-code/src/tui/kimi-tui.ts | 39 +++++++++++++++++++ .../test/tui/commands/registry.test.ts | 1 + .../test/tui/commands/resolve.test.ts | 5 +++ .../agent-core/src/agent/context/index.ts | 28 +++++++++++++ packages/agent-core/src/agent/index.ts | 3 ++ .../agent-core/src/agent/records/index.ts | 3 ++ .../agent-core/src/agent/records/types.ts | 1 + packages/agent-core/src/rpc/core-api.ts | 4 ++ packages/agent-core/src/rpc/core-impl.ts | 5 +++ packages/agent-core/src/session/rpc.ts | 5 +++ packages/node-sdk/src/rpc.ts | 9 +++++ packages/node-sdk/src/session.ts | 5 +++ 14 files changed, 121 insertions(+) diff --git a/apps/kimi-code/src/tui/commands/dispatch.ts b/apps/kimi-code/src/tui/commands/dispatch.ts index 3bd878b0..fafad489 100644 --- a/apps/kimi-code/src/tui/commands/dispatch.ts +++ b/apps/kimi-code/src/tui/commands/dispatch.ts @@ -124,6 +124,9 @@ export interface SlashCommandHost { sendSkillActivation(session: Session, skillName: string, skillArgs: string): void; readonly skillCommandMap: Map; + // Undo + undoLastTurn(): void; + // Controller refs readonly streamingUI: StreamingUIController; readonly tasksBrowserController: TasksBrowserController; @@ -279,6 +282,9 @@ async function handleBuiltInSlashCommand( case 'logout': await handleLogoutCommand(host); return; + case 'undo': + host.undoLastTurn(); + return; default: host.showError(`Unknown slash command: /${String(name)}`); return; diff --git a/apps/kimi-code/src/tui/commands/registry.ts b/apps/kimi-code/src/tui/commands/registry.ts index faf76b57..c2c99cf2 100644 --- a/apps/kimi-code/src/tui/commands/registry.ts +++ b/apps/kimi-code/src/tui/commands/registry.ts @@ -127,6 +127,13 @@ export const BUILTIN_SLASH_COMMANDS = [ priority: 60, availability: 'always', }, + { + name: 'undo', + aliases: [], + description: 'Withdraw the last prompt from the transcript', + priority: 80, + availability: 'idle-only', + }, { name: 'editor', aliases: [], diff --git a/apps/kimi-code/src/tui/kimi-tui.ts b/apps/kimi-code/src/tui/kimi-tui.ts index 9702ed78..1442b939 100644 --- a/apps/kimi-code/src/tui/kimi-tui.ts +++ b/apps/kimi-code/src/tui/kimi-tui.ts @@ -897,6 +897,45 @@ export class KimiTUI { return this.state.transcriptEntries.length > 0; } + undoLastTurn(): void { + if (this.state.appState.streamingPhase !== 'idle') { + this.showError('Cannot undo while streaming — press Esc or Ctrl-C first.'); + return; + } + + const session = this.session; + if (session === undefined) { + this.showError(NO_ACTIVE_SESSION_MESSAGE); + return; + } + + const entries = this.state.transcriptEntries; + const lastUserIndex = entries.findLastIndex((e) => e.kind === 'user'); + if (lastUserIndex < 0) { + this.showError('Nothing to undo.'); + return; + } + + entries.splice(lastUserIndex); + + this.state.transcriptContainer.clear(); + this.clearTerminalInlineImages(); + for (const entry of entries) { + const component = this.createTranscriptComponent(entry); + if (component) { + this.state.transcriptContainer.addChild(component); + } + } + + if (entries.length === 0) { + this.renderWelcome(); + } + + this.state.ui.requestRender(); + + void session.undoHistory(1); + } + async getStartupMcpMs(): Promise { const session = this.session; if (session === undefined) return 0; diff --git a/apps/kimi-code/test/tui/commands/registry.test.ts b/apps/kimi-code/test/tui/commands/registry.test.ts index 74737fb5..c481aa2c 100644 --- a/apps/kimi-code/test/tui/commands/registry.test.ts +++ b/apps/kimi-code/test/tui/commands/registry.test.ts @@ -97,6 +97,7 @@ describe('built-in slash command registry', () => { 'status', 'theme', 'title', + 'undo', 'usage', 'version', 'yolo', diff --git a/apps/kimi-code/test/tui/commands/resolve.test.ts b/apps/kimi-code/test/tui/commands/resolve.test.ts index 07381c0b..c839cba9 100644 --- a/apps/kimi-code/test/tui/commands/resolve.test.ts +++ b/apps/kimi-code/test/tui/commands/resolve.test.ts @@ -63,6 +63,11 @@ describe('resolveSlashCommandInput', () => { commandName: 'resume', reason: 'streaming', }); + expect(resolve('/undo', { isStreaming: true })).toEqual({ + kind: 'blocked', + commandName: 'undo', + reason: 'streaming', + }); }); it('blocks model and session pickers while compacting', () => { diff --git a/packages/agent-core/src/agent/context/index.ts b/packages/agent-core/src/agent/context/index.ts index 4b3810f9..f59cb0a5 100644 --- a/packages/agent-core/src/agent/context/index.ts +++ b/packages/agent-core/src/agent/context/index.ts @@ -71,6 +71,34 @@ export class ContextMemory { this.agent.emitStatusUpdated(); } + undo(count: number): void { + if (count <= 0) return; + if (this._history.length === 0) return; + + this.agent.records.logRecord({ type: 'context.undo', count }); + + let removedUserCount = 0; + for (let i = this._history.length - 1; i >= 0; i--) { + const message = this._history.pop(); + if (message === undefined) continue; + + if (i < this.tokenCountCoveredMessageCount) { + this.tokenCountCoveredMessageCount--; + this._tokenCount -= estimateTokensForMessages([message]); + } + + if (message.role === 'user') { + removedUserCount++; + if (removedUserCount >= count) break; + } + } + + this.openSteps.clear(); + this.pendingToolResultIds.clear(); + this.deferredMessages = []; + this.agent.emitStatusUpdated(); + } + applyCompaction(summary: CompactionResult): void { this.agent.records.logRecord({ type: 'context.apply_compaction', diff --git a/packages/agent-core/src/agent/index.ts b/packages/agent-core/src/agent/index.ts index d38e0c32..a3738789 100644 --- a/packages/agent-core/src/agent/index.ts +++ b/packages/agent-core/src/agent/index.ts @@ -298,6 +298,9 @@ export class Agent { } this.turn.cancel(payload.turnId); }, + undoHistory: (payload) => { + this.context.undo(payload.count); + }, setThinking: (payload) => { const wasEnabled = this.config.thinkingLevel !== 'off'; this.config.update({ thinkingLevel: payload.level }); diff --git a/packages/agent-core/src/agent/records/index.ts b/packages/agent-core/src/agent/records/index.ts index cf79270f..c257c816 100644 --- a/packages/agent-core/src/agent/records/index.ts +++ b/packages/agent-core/src/agent/records/index.ts @@ -82,6 +82,9 @@ function restoreAgentRecord(agent: Agent, input: AgentRecord): void { case 'context.apply_compaction': agent.context.applyCompaction(input); return; + case 'context.undo': + agent.context.undo(input.count); + return; case 'tools.register_user_tool': agent.tools.registerUserTool(input); return; diff --git a/packages/agent-core/src/agent/records/types.ts b/packages/agent-core/src/agent/records/types.ts index c8acefb7..0e78f545 100644 --- a/packages/agent-core/src/agent/records/types.ts +++ b/packages/agent-core/src/agent/records/types.ts @@ -70,6 +70,7 @@ export interface AgentRecordEvents { 'context.append_loop_event': { event: LoopRecordedEvent }; 'context.clear': {}; 'context.apply_compaction': CompactionResult; + 'context.undo': { count: number }; 'tools.update_store': ToolStoreUpdate; } diff --git a/packages/agent-core/src/rpc/core-api.ts b/packages/agent-core/src/rpc/core-api.ts index 504e9a30..ea6be5b9 100644 --- a/packages/agent-core/src/rpc/core-api.ts +++ b/packages/agent-core/src/rpc/core-api.ts @@ -154,6 +154,9 @@ export interface CancelPlanPayload { export interface BeginCompactionPayload { readonly instruction?: string; } +export interface UndoHistoryPayload { + readonly count: number; +} export interface RegisterToolPayload { readonly name: string; readonly description: string; @@ -265,6 +268,7 @@ export interface AgentAPI { prompt: (payload: PromptPayload) => void; steer: (payload: SteerPayload) => void; cancel: (payload: CancelPayload) => void; + undoHistory: (payload: UndoHistoryPayload) => void; setThinking: (payload: SetThinkingPayload) => void; setPermission: (payload: SetPermissionPayload) => void; setModel: (payload: SetModelPayload) => SetModelResult; diff --git a/packages/agent-core/src/rpc/core-impl.ts b/packages/agent-core/src/rpc/core-impl.ts index 26e0f7aa..24d153bf 100644 --- a/packages/agent-core/src/rpc/core-impl.ts +++ b/packages/agent-core/src/rpc/core-impl.ts @@ -84,6 +84,7 @@ import type { SkillSummary, SteerPayload, StopBackgroundPayload, + UndoHistoryPayload, UnregisterToolPayload, UpdateSessionMetadataPayload, } from './core-api'; @@ -430,6 +431,10 @@ export class KimiCore implements PromisableMethods { return this.sessionApi(sessionId).cancel(payload); } + undoHistory({ sessionId, ...payload }: SessionAgentPayload) { + return this.sessionApi(sessionId).undoHistory(payload); + } + async setModel({ sessionId, ...payload diff --git a/packages/agent-core/src/session/rpc.ts b/packages/agent-core/src/session/rpc.ts index be5eac82..883827f5 100644 --- a/packages/agent-core/src/session/rpc.ts +++ b/packages/agent-core/src/session/rpc.ts @@ -23,6 +23,7 @@ import type { SkillSummary, SteerPayload, StopBackgroundPayload, + UndoHistoryPayload, UnregisterToolPayload, UpdateSessionMetadataPayload, } from '#/rpc'; @@ -103,6 +104,10 @@ export class SessionAPIImpl implements PromisableMethods { return this.getAgent(agentId).cancel(payload); } + undoHistory({ agentId, ...payload }: AgentScopedPayload) { + return this.getAgent(agentId).undoHistory(payload); + } + setModel({ agentId, ...payload }: AgentScopedPayload) { return this.getAgent(agentId).setModel(payload); } diff --git a/packages/node-sdk/src/rpc.ts b/packages/node-sdk/src/rpc.ts index 7346e5a5..d66ff679 100644 --- a/packages/node-sdk/src/rpc.ts +++ b/packages/node-sdk/src/rpc.ts @@ -318,6 +318,15 @@ export class SDKRpcClient { }); } + async undoHistory(input: SessionIdRpcInput & { count: number }): Promise { + const rpc = await this.getRpc(); + return rpc.undoHistory({ + sessionId: input.sessionId, + agentId: this.interactiveAgentId, + count: input.count, + }); + } + async getContext(input: SessionIdRpcInput): Promise { const rpc = await this.getRpc(); return rpc.getContext({ diff --git a/packages/node-sdk/src/session.ts b/packages/node-sdk/src/session.ts index 6dc395ef..24d8f286 100644 --- a/packages/node-sdk/src/session.ts +++ b/packages/node-sdk/src/session.ts @@ -166,6 +166,11 @@ export class Session { await this.rpc.cancelCompaction({ sessionId: this.id }); } + async undoHistory(count: number = 1): Promise { + this.ensureOpen(); + await this.rpc.undoHistory({ sessionId: this.id, count }); + } + async getContext(): Promise { this.ensureOpen(); return this.rpc.getContext({ sessionId: this.id }); From 4c373a7819cfb3ddbcac0d9593a61aee879a2729 Mon Sep 17 00:00:00 2001 From: _Kerman Date: Mon, 1 Jun 2026 18:54:50 +0800 Subject: [PATCH 02/10] feat: add /undo slash command and keep replay in sync --- .changeset/feat-undo-command.md | 6 ++ apps/kimi-code/src/tui/kimi-tui.ts | 21 ++-- .../agent-core/src/agent/context/index.ts | 5 + packages/agent-core/src/agent/replay/index.ts | 10 ++ packages/agent-core/test/agent/resume.test.ts | 100 ++++++++++++++++++ 5 files changed, 134 insertions(+), 8 deletions(-) create mode 100644 .changeset/feat-undo-command.md diff --git a/.changeset/feat-undo-command.md b/.changeset/feat-undo-command.md new file mode 100644 index 00000000..b43085f5 --- /dev/null +++ b/.changeset/feat-undo-command.md @@ -0,0 +1,6 @@ +--- +"@moonshot-ai/agent-core": minor +"@moonshot-ai/kimi-code": minor +--- + +Add `/undo` slash command to withdraw the last prompt from conversation history, and keep replay records in sync when a prompt is undone. diff --git a/apps/kimi-code/src/tui/kimi-tui.ts b/apps/kimi-code/src/tui/kimi-tui.ts index 1442b939..336494f9 100644 --- a/apps/kimi-code/src/tui/kimi-tui.ts +++ b/apps/kimi-code/src/tui/kimi-tui.ts @@ -916,17 +916,22 @@ export class KimiTUI { return; } - entries.splice(lastUserIndex); - - this.state.transcriptContainer.clear(); - this.clearTerminalInlineImages(); - for (const entry of entries) { - const component = this.createTranscriptComponent(entry); - if (component) { - this.state.transcriptContainer.addChild(component); + const children = this.state.transcriptContainer.children; + let lastUserComponentIndex = -1; + for (let i = children.length - 1; i >= 0; i--) { + if (children[i] instanceof UserMessageComponent) { + lastUserComponentIndex = i; + break; } } + if (lastUserComponentIndex >= 0) { + children.splice(lastUserComponentIndex); + this.state.transcriptContainer.invalidate(); + } + + entries.splice(lastUserIndex); + if (entries.length === 0) { this.renderWelcome(); } diff --git a/packages/agent-core/src/agent/context/index.ts b/packages/agent-core/src/agent/context/index.ts index f59cb0a5..984f621e 100644 --- a/packages/agent-core/src/agent/context/index.ts +++ b/packages/agent-core/src/agent/context/index.ts @@ -77,11 +77,14 @@ export class ContextMemory { this.agent.records.logRecord({ type: 'context.undo', count }); + let removedCount = 0; let removedUserCount = 0; for (let i = this._history.length - 1; i >= 0; i--) { const message = this._history.pop(); if (message === undefined) continue; + removedCount++; + if (i < this.tokenCountCoveredMessageCount) { this.tokenCountCoveredMessageCount--; this._tokenCount -= estimateTokensForMessages([message]); @@ -93,6 +96,8 @@ export class ContextMemory { } } + this.agent.replayBuilder.removeLastMessages(removedCount); + this.openSteps.clear(); this.pendingToolResultIds.clear(); this.deferredMessages = []; diff --git a/packages/agent-core/src/agent/replay/index.ts b/packages/agent-core/src/agent/replay/index.ts index 0196e622..c9410cf3 100644 --- a/packages/agent-core/src/agent/replay/index.ts +++ b/packages/agent-core/src/agent/replay/index.ts @@ -12,6 +12,16 @@ export class ReplayBuilder { } } + removeLastMessages(count: number): void { + let removed = 0; + for (let i = this.records.length - 1; i >= 0 && removed < count; i--) { + if (this.records[i]!.type === 'message') { + this.records.splice(i, 1); + removed++; + } + } + } + buildResult(): readonly AgentReplayRecord[] { return this.records; } diff --git a/packages/agent-core/test/agent/resume.test.ts b/packages/agent-core/test/agent/resume.test.ts index dafbc864..2d72b39e 100644 --- a/packages/agent-core/test/agent/resume.test.ts +++ b/packages/agent-core/test/agent/resume.test.ts @@ -323,6 +323,106 @@ describe('Agent resume', () => { }), }); }); + + it('removes replay messages matching undone history', async () => { + const persistence = new RecordingAgentPersistence([ + { + type: 'context.append_message', + message: { + role: 'user', + content: [{ type: 'text', text: 'first prompt' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + }, + { + type: 'context.append_loop_event', + event: { + type: 'step.begin', + uuid: 'step-1', + turnId: '0', + step: 1, + }, + }, + { + type: 'context.append_loop_event', + event: { + type: 'content.part', + uuid: 'part-1', + turnId: '0', + step: 1, + stepUuid: 'step-1', + part: { type: 'text', text: 'first response' }, + }, + }, + { + type: 'context.append_loop_event', + event: { + type: 'step.end', + uuid: 'step-1', + turnId: '0', + step: 1, + }, + }, + { + type: 'context.append_message', + message: { + role: 'user', + content: [{ type: 'text', text: 'second prompt' }], + toolCalls: [], + origin: { kind: 'user' }, + }, + }, + { + type: 'context.append_loop_event', + event: { + type: 'step.begin', + uuid: 'step-2', + turnId: '1', + step: 1, + }, + }, + { + type: 'context.append_loop_event', + event: { + type: 'content.part', + uuid: 'part-2', + turnId: '1', + step: 1, + stepUuid: 'step-2', + part: { type: 'text', text: 'second response' }, + }, + }, + { + type: 'context.append_loop_event', + event: { + type: 'step.end', + uuid: 'step-2', + turnId: '1', + step: 1, + }, + }, + { type: 'context.undo', count: 1 }, + ]); + const ctx = testAgent({ persistence }); + + await ctx.agent.resume(); + + expect(ctx.agent.context.history).toHaveLength(2); + expect(ctx.agent.context.history[0]?.role).toBe('user'); + expect(ctx.agent.context.history[1]?.role).toBe('assistant'); + + const replay = ctx.agent.replayBuilder.buildResult(); + expect(replay).toHaveLength(2); + expect(replay[0]).toMatchObject({ + type: 'message', + message: expect.objectContaining({ role: 'user', content: [{ type: 'text', text: 'first prompt' }] }), + }); + expect(replay[1]).toMatchObject({ + type: 'message', + message: expect.objectContaining({ role: 'assistant', content: [{ type: 'text', text: 'first response' }] }), + }); + }); }); class RecordingAgentPersistence extends InMemoryAgentRecordPersistence { From 421adaf23dc3255af1af0151bd6c477e4b16d349 Mon Sep 17 00:00:00 2001 From: _Kerman Date: Mon, 1 Jun 2026 19:18:29 +0800 Subject: [PATCH 03/10] fix: only count real user prompts in undo and include skill-activation turns in TUI undo --- apps/kimi-code/src/tui/kimi-tui.ts | 9 +++-- .../agent-core/src/agent/context/index.ts | 12 ++++++- .../agent-core/test/agent/context.test.ts | 34 +++++++++++++++++++ 3 files changed, 52 insertions(+), 3 deletions(-) diff --git a/apps/kimi-code/src/tui/kimi-tui.ts b/apps/kimi-code/src/tui/kimi-tui.ts index 336494f9..4184ec76 100644 --- a/apps/kimi-code/src/tui/kimi-tui.ts +++ b/apps/kimi-code/src/tui/kimi-tui.ts @@ -910,7 +910,9 @@ export class KimiTUI { } const entries = this.state.transcriptEntries; - const lastUserIndex = entries.findLastIndex((e) => e.kind === 'user'); + const lastUserIndex = entries.findLastIndex( + (e) => e.kind === 'user' || e.kind === 'skill_activation', + ); if (lastUserIndex < 0) { this.showError('Nothing to undo.'); return; @@ -919,7 +921,10 @@ export class KimiTUI { const children = this.state.transcriptContainer.children; let lastUserComponentIndex = -1; for (let i = children.length - 1; i >= 0; i--) { - if (children[i] instanceof UserMessageComponent) { + if ( + children[i] instanceof UserMessageComponent || + children[i] instanceof SkillActivationComponent + ) { lastUserComponentIndex = i; break; } diff --git a/packages/agent-core/src/agent/context/index.ts b/packages/agent-core/src/agent/context/index.ts index 984f621e..a17a7bf9 100644 --- a/packages/agent-core/src/agent/context/index.ts +++ b/packages/agent-core/src/agent/context/index.ts @@ -90,7 +90,7 @@ export class ContextMemory { this._tokenCount -= estimateTokensForMessages([message]); } - if (message.role === 'user') { + if (isRealUserPrompt(message)) { removedUserCount++; if (removedUserCount >= count) break; } @@ -296,3 +296,13 @@ function toolResultOutputForModel(result: ExecutableToolResult): string | Conten function isEmptyOutputText(output: string): boolean { return output.length === 0 || output.trim() === TOOL_OUTPUT_EMPTY_TEXT; } + +function isRealUserPrompt(message: ContextMessage): boolean { + if (message.role !== 'user') return false; + const origin = message.origin; + if (origin === undefined || origin.kind === 'user') return true; + if (origin.kind === 'skill_activation') { + return origin.trigger === 'user-slash'; + } + return false; +} diff --git a/packages/agent-core/test/agent/context.test.ts b/packages/agent-core/test/agent/context.test.ts index 5bf02e9f..5f5f45ea 100644 --- a/packages/agent-core/test/agent/context.test.ts +++ b/packages/agent-core/test/agent/context.test.ts @@ -507,6 +507,40 @@ describe('Agent context', () => { ); }); + it('undo only counts real user prompts, skipping background notifications', () => { + const ctx = testAgent(); + ctx.configure(); + + ctx.appendAssistantText(1, 'first response'); + ctx.appendAssistantText(2, 'second response'); + + // Append a background task notification (role: 'user' but not a real 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', + }, + }); + + expect(ctx.agent.context.history.map((m) => m.role)).toEqual([ + 'user', + 'assistant', + 'user', + 'assistant', + 'user', + ]); + + ctx.agent.context.undo(1); + + // Should remove the background notification, the second assistant, and the second user prompt + expect(ctx.agent.context.history.map((m) => m.role)).toEqual(['user', 'assistant']); + }); + }); describe('Agent context notification projection', () => { From 5e0675eac7f04dfd5306bd34713f04de7bfef2f9 Mon Sep 17 00:00:00 2001 From: _Kerman Date: Mon, 1 Jun 2026 20:40:32 +0800 Subject: [PATCH 04/10] fix: keep undo state consistent --- apps/kimi-code/src/tui/commands/dispatch.ts | 4 +- apps/kimi-code/src/tui/kimi-tui.ts | 12 ++++-- .../test/tui/kimi-tui-message-flow.test.ts | 27 ++++++++++++ .../agent-core/src/agent/context/index.ts | 15 +++++-- .../src/agent/injection/injector.ts | 9 ++++ .../agent-core/src/agent/injection/manager.ts | 6 +++ packages/agent-core/src/agent/replay/index.ts | 11 ++--- .../agent-core/test/agent/context.test.ts | 42 +++++++++++++++++++ packages/agent-core/test/agent/plan.test.ts | 16 +++++++ 9 files changed, 128 insertions(+), 14 deletions(-) diff --git a/apps/kimi-code/src/tui/commands/dispatch.ts b/apps/kimi-code/src/tui/commands/dispatch.ts index fafad489..293d0aae 100644 --- a/apps/kimi-code/src/tui/commands/dispatch.ts +++ b/apps/kimi-code/src/tui/commands/dispatch.ts @@ -125,7 +125,7 @@ export interface SlashCommandHost { readonly skillCommandMap: Map; // Undo - undoLastTurn(): void; + undoLastTurn(): Promise; // Controller refs readonly streamingUI: StreamingUIController; @@ -283,7 +283,7 @@ async function handleBuiltInSlashCommand( await handleLogoutCommand(host); return; case 'undo': - host.undoLastTurn(); + await host.undoLastTurn(); return; default: host.showError(`Unknown slash command: /${String(name)}`); diff --git a/apps/kimi-code/src/tui/kimi-tui.ts b/apps/kimi-code/src/tui/kimi-tui.ts index 4184ec76..4fcd1e87 100644 --- a/apps/kimi-code/src/tui/kimi-tui.ts +++ b/apps/kimi-code/src/tui/kimi-tui.ts @@ -897,7 +897,7 @@ export class KimiTUI { return this.state.transcriptEntries.length > 0; } - undoLastTurn(): void { + async undoLastTurn(): Promise { if (this.state.appState.streamingPhase !== 'idle') { this.showError('Cannot undo while streaming — press Esc or Ctrl-C first.'); return; @@ -918,6 +918,14 @@ export class KimiTUI { return; } + try { + await session.undoHistory(1); + } catch (error) { + const message = formatErrorMessage(error); + this.showError(`Failed to undo: ${message}`); + return; + } + const children = this.state.transcriptContainer.children; let lastUserComponentIndex = -1; for (let i = children.length - 1; i >= 0; i--) { @@ -942,8 +950,6 @@ export class KimiTUI { } this.state.ui.requestRender(); - - void session.undoHistory(1); } async getStartupMcpMs(): Promise { diff --git a/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts b/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts index 58c6f15a..3b91b0ab 100644 --- a/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts +++ b/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts @@ -56,6 +56,7 @@ interface MessageDriver { handleUserInput(text: string): void; persistInputHistory(text: string): Promise; getCurrentSessionId(): string; + undoLastTurn(): Promise; } interface FeedbackDriver extends MessageDriver { @@ -110,6 +111,7 @@ function makeSession(overrides: Record = {}) { prompt: vi.fn(async () => {}), steer: vi.fn(async () => {}), init: vi.fn(async () => {}), + undoHistory: vi.fn(async () => {}), cancel: vi.fn(async () => {}), cancelCompaction: vi.fn(async () => {}), getStatus: vi.fn(async () => ({ @@ -685,6 +687,31 @@ describe('KimiTUI message flow', () => { ]); }); + it('keeps the transcript intact when undo RPC fails', async () => { + const session = makeSession({ + undoHistory: vi.fn(async () => { + throw new Error('core rpc unavailable'); + }), + }); + const { driver } = await makeDriver(session); + + driver.handleUserInput('hello'); + driver.state.appState.streamingPhase = 'idle'; + + await driver.undoLastTurn(); + + expect(session.undoHistory).toHaveBeenCalledWith(1); + expect(driver.state.transcriptEntries).toEqual([ + expect.objectContaining({ + kind: 'user', + content: 'hello', + }), + ]); + const transcript = stripSgr(renderTranscript(driver)); + expect(transcript).toContain('hello'); + expect(transcript).toContain('Error: Failed to undo: core rpc unavailable'); + }); + it('sends pasted image placeholders as image content parts', async () => { const { driver, session } = await makeDriver(); const imageStore = (driver as unknown as { imageStore: ImageAttachmentStore }).imageStore; diff --git a/packages/agent-core/src/agent/context/index.ts b/packages/agent-core/src/agent/context/index.ts index a17a7bf9..81c1b7ad 100644 --- a/packages/agent-core/src/agent/context/index.ts +++ b/packages/agent-core/src/agent/context/index.ts @@ -77,13 +77,16 @@ export class ContextMemory { this.agent.records.logRecord({ type: 'context.undo', count }); - let removedCount = 0; let removedUserCount = 0; + const removedMessages = new Set(); for (let i = this._history.length - 1; i >= 0; i--) { - const message = this._history.pop(); + const message = this._history[i]; if (message === undefined) continue; + if (isInjectionMessage(message)) continue; - removedCount++; + removedMessages.add(message); + this._history.splice(i, 1); + this.agent.injection.onContextMessageRemoved(i); if (i < this.tokenCountCoveredMessageCount) { this.tokenCountCoveredMessageCount--; @@ -96,7 +99,7 @@ export class ContextMemory { } } - this.agent.replayBuilder.removeLastMessages(removedCount); + this.agent.replayBuilder.removeLastMessages(removedMessages); this.openSteps.clear(); this.pendingToolResultIds.clear(); @@ -306,3 +309,7 @@ function isRealUserPrompt(message: ContextMessage): boolean { } return false; } + +function isInjectionMessage(message: ContextMessage): boolean { + return message.origin?.kind === 'injection'; +} diff --git a/packages/agent-core/src/agent/injection/injector.ts b/packages/agent-core/src/agent/injection/injector.ts index 1084e32e..504e412d 100644 --- a/packages/agent-core/src/agent/injection/injector.ts +++ b/packages/agent-core/src/agent/injection/injector.ts @@ -16,6 +16,15 @@ export abstract class DynamicInjector { } } + onContextMessageRemoved(index: number): void { + if (this.injectedAt === null) return; + if (index < this.injectedAt) { + this.injectedAt--; + } else if (index === this.injectedAt) { + this.injectedAt = null; + } + } + async inject(): Promise { const injection = await this.getInjection(); if (injection) { diff --git a/packages/agent-core/src/agent/injection/manager.ts b/packages/agent-core/src/agent/injection/manager.ts index edda42c4..1de6aa92 100644 --- a/packages/agent-core/src/agent/injection/manager.ts +++ b/packages/agent-core/src/agent/injection/manager.ts @@ -36,4 +36,10 @@ export class InjectionManager { } } } + + onContextMessageRemoved(index: number): void { + for (const injector of this.injectors) { + injector.onContextMessageRemoved(index); + } + } } diff --git a/packages/agent-core/src/agent/replay/index.ts b/packages/agent-core/src/agent/replay/index.ts index c9410cf3..acb76005 100644 --- a/packages/agent-core/src/agent/replay/index.ts +++ b/packages/agent-core/src/agent/replay/index.ts @@ -1,5 +1,6 @@ import type { Agent } from '..'; import type { AgentReplayRecord } from '../..'; +import type { ContextMessage } from '../context'; export class ReplayBuilder { protected readonly records: AgentReplayRecord[] = []; @@ -12,12 +13,12 @@ export class ReplayBuilder { } } - removeLastMessages(count: number): void { - let removed = 0; - for (let i = this.records.length - 1; i >= 0 && removed < count; i--) { - if (this.records[i]!.type === 'message') { + removeLastMessages(removedMessages: ReadonlySet): void { + if (removedMessages.size === 0) return; + for (let i = this.records.length - 1; i >= 0; i--) { + const record = this.records[i]!; + if (record.type === 'message' && removedMessages.has(record.message)) { this.records.splice(i, 1); - removed++; } } } diff --git a/packages/agent-core/test/agent/context.test.ts b/packages/agent-core/test/agent/context.test.ts index 5f5f45ea..e5953c64 100644 --- a/packages/agent-core/test/agent/context.test.ts +++ b/packages/agent-core/test/agent/context.test.ts @@ -541,6 +541,48 @@ describe('Agent context', () => { expect(ctx.agent.context.history.map((m) => m.role)).toEqual(['user', 'assistant']); }); + it('preserves injection messages when undo removes the surrounding turn', () => { + const ctx = testAgent(); + ctx.configure(); + + ctx.dispatch({ + type: 'context.append_message', + message: userMessage('do the work', { kind: 'user' }), + }); + ctx.dispatch({ + type: 'context.append_message', + message: userMessage('Plan mode is active', { + kind: 'injection', + variant: 'plan_mode', + }), + }); + ctx.dispatch({ + type: 'context.append_message', + message: { + role: 'assistant', + content: [{ type: 'text', text: 'work done' }], + toolCalls: [], + }, + }); + + ctx.agent.context.undo(1); + + expect(ctx.agent.context.history).toEqual([ + expect.objectContaining({ + role: 'user', + origin: { kind: 'injection', variant: 'plan_mode' }, + }), + ]); + expect(ctx.agent.replayBuilder.buildResult()).toEqual([ + expect.objectContaining({ + type: 'message', + message: expect.objectContaining({ + origin: { kind: 'injection', variant: 'plan_mode' }, + }), + }), + ]); + }); + }); describe('Agent context notification projection', () => { diff --git a/packages/agent-core/test/agent/plan.test.ts b/packages/agent-core/test/agent/plan.test.ts index 8c32d48d..2f652584 100644 --- a/packages/agent-core/test/agent/plan.test.ts +++ b/packages/agent-core/test/agent/plan.test.ts @@ -579,6 +579,22 @@ describe('plan mode injection cadence', () => { expect(ctx.agent.context.history).toHaveLength(afterExit); await ctx.expectResumeMatches(); }); + + it('keeps the preserved injection index aligned after undo removes earlier messages', async () => { + const ctx = testAgent(); + ctx.configure(); + await ctx.agent.planMode.enter('test-plan', false); + + ctx.agent.context.appendUserMessage([{ type: 'text', text: 'draft the plan' }]); + await ctx.agent.injection.inject(); + ctx.appendAssistantTurn(1, 'Plan drafted.'); + + ctx.agent.context.undo(1); + ctx.agent.context.appendUserMessage([{ type: 'text', text: 'new plan request' }]); + await ctx.agent.injection.inject(); + + expect(lastUserText(ctx.agent.context.history)).toContain('Plan mode is active'); + }); }); function delay(ms: number): Promise { From 94291833414a15ba0246dc03b040859d01b7e507 Mon Sep 17 00:00:00 2001 From: _Kerman Date: Mon, 1 Jun 2026 22:13:30 +0800 Subject: [PATCH 05/10] fix --- apps/kimi-code/src/tui/kimi-tui.ts | 7 ++ .../test/tui/kimi-tui-message-flow.test.ts | 16 +++++ .../agent-core/src/agent/compaction/micro.ts | 4 +- .../agent-core/src/agent/context/index.ts | 23 +++++-- .../test/agent/compaction/micro.test.ts | 33 +++++++++ .../agent-core/test/agent/context.test.ts | 67 +++++++++++++++++++ 6 files changed, 143 insertions(+), 7 deletions(-) diff --git a/apps/kimi-code/src/tui/kimi-tui.ts b/apps/kimi-code/src/tui/kimi-tui.ts index 4fcd1e87..190a5c5c 100644 --- a/apps/kimi-code/src/tui/kimi-tui.ts +++ b/apps/kimi-code/src/tui/kimi-tui.ts @@ -1328,6 +1328,13 @@ export class KimiTUI { } private renderWelcome(): void { + if ( + this.state.transcriptContainer.children.some( + (child) => child instanceof WelcomeComponent, + ) + ) { + return; + } const welcome = new WelcomeComponent(this.state.appState, this.state.theme.colors); this.state.transcriptContainer.addChild(welcome); } diff --git a/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts b/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts index 3b91b0ab..3db951ed 100644 --- a/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts +++ b/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts @@ -12,6 +12,7 @@ import { afterEach, describe, expect, it, vi } from 'vitest'; import { ApprovalPanelComponent } from '#/tui/components/dialogs/approval-panel'; import { KIMI_CODE_PLUGIN_MARKETPLACE_URL } from '#/constant/app'; +import { WelcomeComponent } from '#/tui/components/chrome/welcome'; import { ModelSelectorComponent } from '#/tui/components/dialogs/model-selector'; import { PluginMcpSelectorComponent, @@ -712,6 +713,21 @@ describe('KimiTUI message flow', () => { expect(transcript).toContain('Error: Failed to undo: core rpc unavailable'); }); + it('does not duplicate welcome after undoing the only turn', async () => { + const { driver } = await makeDriver(); + + driver.handleUserInput('hello'); + driver.state.appState.streamingPhase = 'idle'; + + await driver.undoLastTurn(); + + expect( + driver.state.transcriptContainer.children.filter( + (child) => child instanceof WelcomeComponent, + ), + ).toHaveLength(1); + }); + it('sends pasted image placeholders as image content parts', async () => { const { driver, session } = await makeDriver(); const imageStore = (driver as unknown as { imageStore: ImageAttachmentStore }).imageStore; diff --git a/packages/agent-core/src/agent/compaction/micro.ts b/packages/agent-core/src/agent/compaction/micro.ts index 938f8769..2dbfa5a5 100644 --- a/packages/agent-core/src/agent/compaction/micro.ts +++ b/packages/agent-core/src/agent/compaction/micro.ts @@ -32,8 +32,8 @@ export class MicroCompaction { this.config = { ...DEFAULT_CONFIG, ...config }; } - reset(): void { - this.cutoff = 0; + reset(maxCutoff = 0): void { + this.cutoff = Math.min(this.cutoff, maxCutoff); } apply(cutoff: number): void { diff --git a/packages/agent-core/src/agent/context/index.ts b/packages/agent-core/src/agent/context/index.ts index 81c1b7ad..386dccd9 100644 --- a/packages/agent-core/src/agent/context/index.ts +++ b/packages/agent-core/src/agent/context/index.ts @@ -1,6 +1,7 @@ import { createToolMessage, type ContentPart, type Message } from '@moonshot-ai/kosong'; import type { Agent } from '..'; +import { ErrorCodes, KimiError } from '../../errors'; import type { ExecutableToolResult, LoopRecordedEvent } from '../../loop'; import { estimateTokensForMessages } from '../../utils/tokens'; import type { CompactionResult } from '../compaction'; @@ -79,10 +80,15 @@ export class ContextMemory { let removedUserCount = 0; const removedMessages = new Set(); + let stoppedAtBoundary = false; for (let i = this._history.length - 1; i >= 0; i--) { const message = this._history[i]; if (message === undefined) continue; - if (isInjectionMessage(message)) continue; + if (message.origin?.kind === 'injection') continue; + if (message.origin?.kind === 'compaction_summary') { + stoppedAtBoundary = true; + break; + } removedMessages.add(message); this._history.splice(i, 1); @@ -104,7 +110,18 @@ export class ContextMemory { this.openSteps.clear(); this.pendingToolResultIds.clear(); this.deferredMessages = []; + this.agent.microCompaction.reset(this._history.length); this.agent.emitStatusUpdated(); + + if ( + !this.agent.records.restoring && + (stoppedAtBoundary || removedUserCount < count) + ) { + throw new KimiError( + ErrorCodes.REQUEST_INVALID, + 'Nothing to undo in the active context.', + ); + } } applyCompaction(summary: CompactionResult): void { @@ -309,7 +326,3 @@ function isRealUserPrompt(message: ContextMessage): boolean { } return false; } - -function isInjectionMessage(message: ContextMessage): boolean { - return message.origin?.kind === 'injection'; -} diff --git a/packages/agent-core/test/agent/compaction/micro.test.ts b/packages/agent-core/test/agent/compaction/micro.test.ts index c0737954..7975f6a0 100644 --- a/packages/agent-core/test/agent/compaction/micro.test.ts +++ b/packages/agent-core/test/agent/compaction/micro.test.ts @@ -405,6 +405,39 @@ describe('MicroCompaction', () => { ]); }); + it('clamps cutoff when undo shortens the context', () => { + vi.useFakeTimers(); + const ctx = testAgent({ + microCompaction: { + keepRecentMessages: 2, + minContentTokens: 1, + cacheMissedThresholdMs: 60 * MINUTE, + minContextUsageRatio: 0, + }, + }); + + vi.setSystemTime(0); + appendMicroToolExchange(ctx, 1, { output: 'result one' }); + appendMicroToolExchange(ctx, 2, { output: 'result two' }); + appendMicroToolExchange(ctx, 3, { output: 'result three' }); + + vi.setSystemTime(61 * MINUTE); + ctx.agent.microCompaction.detect(); + expect(toolTexts(ctx.agent.context.messages)).toEqual([ + DEFAULT_MARKER, + DEFAULT_MARKER, + 'result three', + ]); + + ctx.agent.context.undo(2); + appendMicroToolExchange(ctx, 4, { output: 'result four' }); + + expect(toolTexts(ctx.agent.context.messages)).toEqual([ + DEFAULT_MARKER, + 'result four', + ]); + }); + it('tracks telemetry when a cache miss advances the micro-compaction cutoff', () => { vi.useFakeTimers(); const records: TelemetryRecord[] = []; diff --git a/packages/agent-core/test/agent/context.test.ts b/packages/agent-core/test/agent/context.test.ts index e5953c64..610b97ca 100644 --- a/packages/agent-core/test/agent/context.test.ts +++ b/packages/agent-core/test/agent/context.test.ts @@ -541,6 +541,73 @@ describe('Agent context', () => { expect(ctx.agent.context.history.map((m) => m.role)).toEqual(['user', 'assistant']); }); + it('stops at compaction summary and records the requested undo count', () => { + const ctx = testAgent(); + ctx.configure(); + ctx.agent.context.appendUserMessage([{ type: 'text', text: 'old user message' }]); + ctx.agent.context.applyCompaction({ + summary: 'summary of compacted context', + compactedCount: 1, + tokensBefore: 100, + tokensAfter: 20, + }); + ctx.agent.context.appendUserMessage([{ type: 'text', text: 'recent user message' }]); + ctx.agent.context.appendMessage({ + role: 'assistant', + content: [{ type: 'text', text: 'recent answer' }], + toolCalls: [], + }); + ctx.newEvents(); + + expect(() => { + ctx.agent.context.undo(2); + }).toThrow('Nothing to undo in the active context.'); + + expect(ctx.agent.context.history).toEqual([ + expect.objectContaining({ + role: 'assistant', + origin: { kind: 'compaction_summary' }, + content: [{ type: 'text', text: 'summary of compacted context' }], + }), + ]); + expect(ctx.newEvents()).toContainEqual( + expect.objectContaining({ + type: '[wire]', + event: 'context.undo', + args: expect.objectContaining({ count: 2 }), + }), + ); + }); + + it('does not throw while restoring an undo that stops at compaction summary', () => { + const ctx = testAgent(); + ctx.configure(); + ctx.agent.context.appendUserMessage([{ type: 'text', text: 'old user message' }]); + ctx.agent.context.applyCompaction({ + summary: 'summary of compacted context', + compactedCount: 1, + tokensBefore: 100, + tokensAfter: 20, + }); + ctx.agent.context.appendUserMessage([{ type: 'text', text: 'recent user message' }]); + ctx.agent.context.appendMessage({ + role: 'assistant', + content: [{ type: 'text', text: 'recent answer' }], + toolCalls: [], + }); + + expect(() => { + ctx.agent.records.restore({ type: 'context.undo', count: 2 }); + }).not.toThrow(); + expect(ctx.agent.context.history).toEqual([ + expect.objectContaining({ + role: 'assistant', + origin: { kind: 'compaction_summary' }, + content: [{ type: 'text', text: 'summary of compacted context' }], + }), + ]); + }); + it('preserves injection messages when undo removes the surrounding turn', () => { const ctx = testAgent(); ctx.configure(); From e296ec1dadacca6ebba03912e620d71df1814abe Mon Sep 17 00:00:00 2001 From: _Kerman Date: Mon, 1 Jun 2026 22:26:35 +0800 Subject: [PATCH 06/10] fix: align tui undo skill anchors --- .../components/messages/skill-activation.ts | 8 ++- .../tui/controllers/session-event-handler.ts | 1 + .../src/tui/controllers/session-replay.ts | 1 + apps/kimi-code/src/tui/kimi-tui.ts | 17 ++++-- apps/kimi-code/src/tui/types.ts | 3 ++ .../kimi-code/src/tui/utils/message-replay.ts | 3 ++ .../test/tui/kimi-tui-message-flow.test.ts | 53 +++++++++++++++++++ 7 files changed, 80 insertions(+), 6 deletions(-) diff --git a/apps/kimi-code/src/tui/components/messages/skill-activation.ts b/apps/kimi-code/src/tui/components/messages/skill-activation.ts index 7f9e05c9..4526328c 100644 --- a/apps/kimi-code/src/tui/components/messages/skill-activation.ts +++ b/apps/kimi-code/src/tui/components/messages/skill-activation.ts @@ -16,11 +16,17 @@ import { Container, Text, Spacer } from '@earendil-works/pi-tui'; import chalk from 'chalk'; import type { ColorPalette } from '#/tui/theme/colors'; +import type { SkillActivationTrigger } from '#/tui/types'; const ARGS_PREVIEW_MAX = 200; export class SkillActivationComponent extends Container { - constructor(name: string, args: string | undefined, colors: ColorPalette) { + constructor( + name: string, + args: string | undefined, + colors: ColorPalette, + readonly trigger?: SkillActivationTrigger, + ) { super(); this.addChild(new Spacer(1)); const head = diff --git a/apps/kimi-code/src/tui/controllers/session-event-handler.ts b/apps/kimi-code/src/tui/controllers/session-event-handler.ts index 27b13a13..0a961c84 100644 --- a/apps/kimi-code/src/tui/controllers/session-event-handler.ts +++ b/apps/kimi-code/src/tui/controllers/session-event-handler.ts @@ -662,6 +662,7 @@ export class SessionEventHandler { skillActivationId: event.activationId, skillName: event.skillName, skillArgs: event.skillArgs, + skillTrigger: event.trigger, }); } diff --git a/apps/kimi-code/src/tui/controllers/session-replay.ts b/apps/kimi-code/src/tui/controllers/session-replay.ts index 27b76ceb..fa1af75f 100644 --- a/apps/kimi-code/src/tui/controllers/session-replay.ts +++ b/apps/kimi-code/src/tui/controllers/session-replay.ts @@ -345,6 +345,7 @@ export class SessionReplayRenderer { skillActivationId: skill.activationId, skillName: skill.skillName, skillArgs: skill.skillArgs, + skillTrigger: skill.trigger, }); } diff --git a/apps/kimi-code/src/tui/kimi-tui.ts b/apps/kimi-code/src/tui/kimi-tui.ts index 190a5c5c..b37f147d 100644 --- a/apps/kimi-code/src/tui/kimi-tui.ts +++ b/apps/kimi-code/src/tui/kimi-tui.ts @@ -910,9 +910,7 @@ export class KimiTUI { } const entries = this.state.transcriptEntries; - const lastUserIndex = entries.findLastIndex( - (e) => e.kind === 'user' || e.kind === 'skill_activation', - ); + const lastUserIndex = entries.findLastIndex(isUndoAnchorEntry); if (lastUserIndex < 0) { this.showError('Nothing to undo.'); return; @@ -929,9 +927,10 @@ export class KimiTUI { const children = this.state.transcriptContainer.children; let lastUserComponentIndex = -1; for (let i = children.length - 1; i >= 0; i--) { + const child = children[i]; if ( - children[i] instanceof UserMessageComponent || - children[i] instanceof SkillActivationComponent + child instanceof UserMessageComponent || + (child instanceof SkillActivationComponent && child.trigger === 'user-slash') ) { lastUserComponentIndex = i; break; @@ -1232,6 +1231,7 @@ export class KimiTUI { entry.skillName ?? entry.content, entry.skillArgs, this.state.theme.colors, + entry.skillTrigger, ); case 'cron': return new CronMessageComponent( @@ -1835,3 +1835,10 @@ export class KimiTUI { } } + +function isUndoAnchorEntry(entry: TranscriptEntry): boolean { + return ( + entry.kind === 'user' || + (entry.kind === 'skill_activation' && entry.skillTrigger === 'user-slash') + ); +} diff --git a/apps/kimi-code/src/tui/types.ts b/apps/kimi-code/src/tui/types.ts index 5a1a8ee4..cc5526c7 100644 --- a/apps/kimi-code/src/tui/types.ts +++ b/apps/kimi-code/src/tui/types.ts @@ -114,6 +114,8 @@ export type TranscriptEntryKind = | 'skill_activation' | 'cron'; +export type SkillActivationTrigger = 'user-slash' | 'model-tool' | 'nested-skill'; + export interface TranscriptEntry { id: string; kind: TranscriptEntryKind; @@ -130,6 +132,7 @@ export interface TranscriptEntry { skillActivationId?: string; skillName?: string; skillArgs?: string; + skillTrigger?: SkillActivationTrigger; } export type LivePaneMode = diff --git a/apps/kimi-code/src/tui/utils/message-replay.ts b/apps/kimi-code/src/tui/utils/message-replay.ts index 8b83186d..3f1638eb 100644 --- a/apps/kimi-code/src/tui/utils/message-replay.ts +++ b/apps/kimi-code/src/tui/utils/message-replay.ts @@ -11,6 +11,7 @@ import type { import type { AppState, BackgroundAgentMetadata, + SkillActivationTrigger, ToolCallBlockData, TranscriptEntry, } from '#/tui/types'; @@ -38,6 +39,7 @@ export interface SkillActivationProjection { readonly activationId: string; readonly skillName: string; readonly skillArgs?: string; + readonly trigger: SkillActivationTrigger; } export interface ReplayBackgroundProjection { @@ -203,6 +205,7 @@ export function skillActivationFromOrigin( activationId: origin.activationId, skillName: origin.skillName, skillArgs: origin.skillArgs, + trigger: origin.trigger, }; } diff --git a/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts b/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts index 3db951ed..42bf0ca9 100644 --- a/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts +++ b/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts @@ -728,6 +728,59 @@ describe('KimiTUI message flow', () => { ).toHaveLength(1); }); + it('undoes from the real user turn when the last skill activation came from the model', async () => { + const { driver } = await makeDriver(); + + driver.handleUserInput('hello'); + driver.sessionEventHandler.handleEvent( + { + type: 'skill.activated', + agentId: 'main', + activationId: 'act-model', + skillName: 'review', + trigger: 'model-tool', + } as Event, + () => {}, + ); + driver.state.appState.streamingPhase = 'idle'; + + await driver.undoLastTurn(); + + expect(driver.state.transcriptEntries).toEqual([]); + const transcript = stripSgr(renderTranscript(driver)); + expect(transcript).not.toContain('hello'); + expect(transcript).not.toContain('review'); + }); + + it('keeps user-slash skill activations as undo anchors', async () => { + const { driver } = await makeDriver(); + + driver.handleUserInput('hello'); + driver.sessionEventHandler.handleEvent( + { + type: 'skill.activated', + agentId: 'main', + activationId: 'act-user', + skillName: 'review', + trigger: 'user-slash', + } as Event, + () => {}, + ); + driver.state.appState.streamingPhase = 'idle'; + + await driver.undoLastTurn(); + + expect(driver.state.transcriptEntries).toEqual([ + expect.objectContaining({ + kind: 'user', + content: 'hello', + }), + ]); + const transcript = stripSgr(renderTranscript(driver)); + expect(transcript).toContain('hello'); + expect(transcript).not.toContain('review'); + }); + it('sends pasted image placeholders as image content parts', async () => { const { driver, session } = await makeDriver(); const imageStore = (driver as unknown as { imageStore: ImageAttachmentStore }).imageStore; From b79ffb04315b5d0a8b22a91b400893e0055fd415 Mon Sep 17 00:00:00 2001 From: _Kerman Date: Tue, 2 Jun 2026 11:22:56 +0800 Subject: [PATCH 07/10] fix --- apps/kimi-code/src/tui/commands/dispatch.ts | 7 +- apps/kimi-code/src/tui/commands/index.ts | 1 + apps/kimi-code/src/tui/commands/undo.ts | 85 +++++++++++++++++++ apps/kimi-code/src/tui/kimi-tui.ts | 61 ------------- .../test/tui/kimi-tui-message-flow.test.ts | 37 ++++++-- 5 files changed, 119 insertions(+), 72 deletions(-) create mode 100644 apps/kimi-code/src/tui/commands/undo.ts diff --git a/apps/kimi-code/src/tui/commands/dispatch.ts b/apps/kimi-code/src/tui/commands/dispatch.ts index 293d0aae..993aeb61 100644 --- a/apps/kimi-code/src/tui/commands/dispatch.ts +++ b/apps/kimi-code/src/tui/commands/dispatch.ts @@ -42,6 +42,7 @@ import { handleInitCommand, handleTitleCommand, } from './session'; +import { handleUndoCommand } from './undo'; // --------------------------------------------------------------------------- // Re-exports — keep existing consumers working @@ -78,6 +79,7 @@ export { handleInitCommand, handleTitleCommand, } from './session'; +export { handleUndoCommand } from './undo'; // --------------------------------------------------------------------------- // Host interface @@ -124,9 +126,6 @@ export interface SlashCommandHost { sendSkillActivation(session: Session, skillName: string, skillArgs: string): void; readonly skillCommandMap: Map; - // Undo - undoLastTurn(): Promise; - // Controller refs readonly streamingUI: StreamingUIController; readonly tasksBrowserController: TasksBrowserController; @@ -283,7 +282,7 @@ async function handleBuiltInSlashCommand( await handleLogoutCommand(host); return; case 'undo': - await host.undoLastTurn(); + await handleUndoCommand(host); return; default: host.showError(`Unknown slash command: /${String(name)}`); diff --git a/apps/kimi-code/src/tui/commands/index.ts b/apps/kimi-code/src/tui/commands/index.ts index 60178b26..261f70ae 100644 --- a/apps/kimi-code/src/tui/commands/index.ts +++ b/apps/kimi-code/src/tui/commands/index.ts @@ -34,6 +34,7 @@ export { handleInitCommand, handleTitleCommand, } from './session'; +export { handleUndoCommand } from './undo'; export { promptApiKey, promptCatalogProviderSelection, diff --git a/apps/kimi-code/src/tui/commands/undo.ts b/apps/kimi-code/src/tui/commands/undo.ts new file mode 100644 index 00000000..46dbb6c3 --- /dev/null +++ b/apps/kimi-code/src/tui/commands/undo.ts @@ -0,0 +1,85 @@ +import { WelcomeComponent } from '../components/chrome/welcome'; +import { SkillActivationComponent } from '../components/messages/skill-activation'; +import { UserMessageComponent } from '../components/messages/user-message'; +import { NO_ACTIVE_SESSION_MESSAGE } from '../constant/kimi-tui'; +import type { TranscriptEntry } from '../types'; +import { formatErrorMessage } from '../utils/event-payload'; +import type { SlashCommandHost } from './dispatch'; + +// --------------------------------------------------------------------------- +// Undo command +// --------------------------------------------------------------------------- + +export async function handleUndoCommand(host: SlashCommandHost): Promise { + if (host.state.appState.streamingPhase !== 'idle') { + host.showError('Cannot undo while streaming — press Esc or Ctrl-C first.'); + return; + } + + const session = host.session; + if (session === undefined) { + host.showError(NO_ACTIVE_SESSION_MESSAGE); + return; + } + + const entries = host.state.transcriptEntries; + const lastUserIndex = entries.findLastIndex(isUndoAnchorEntry); + if (lastUserIndex < 0) { + host.showError('Nothing to undo.'); + return; + } + + try { + await session.undoHistory(1); + } catch (error) { + const message = formatErrorMessage(error); + host.showError(`Failed to undo: ${message}`); + return; + } + + const children = host.state.transcriptContainer.children; + let lastUserComponentIndex = -1; + for (let i = children.length - 1; i >= 0; i--) { + const child = children[i]; + if ( + child instanceof UserMessageComponent || + (child instanceof SkillActivationComponent && child.trigger === 'user-slash') + ) { + lastUserComponentIndex = i; + break; + } + } + + if (lastUserComponentIndex >= 0) { + children.splice(lastUserComponentIndex); + host.state.transcriptContainer.invalidate(); + } + + entries.splice(lastUserIndex); + + if (entries.length === 0) { + renderWelcome(host); + } + + host.state.ui.requestRender(); +} + +function isUndoAnchorEntry(entry: TranscriptEntry): boolean { + return ( + entry.kind === 'user' || + (entry.kind === 'skill_activation' && entry.skillTrigger === 'user-slash') + ); +} + +function renderWelcome(host: SlashCommandHost): void { + if ( + host.state.transcriptContainer.children.some( + (child) => child instanceof WelcomeComponent, + ) + ) { + return; + } + host.state.transcriptContainer.addChild( + new WelcomeComponent(host.state.appState, host.state.theme.colors), + ); +} diff --git a/apps/kimi-code/src/tui/kimi-tui.ts b/apps/kimi-code/src/tui/kimi-tui.ts index b37f147d..66458ed0 100644 --- a/apps/kimi-code/src/tui/kimi-tui.ts +++ b/apps/kimi-code/src/tui/kimi-tui.ts @@ -897,60 +897,6 @@ export class KimiTUI { return this.state.transcriptEntries.length > 0; } - async undoLastTurn(): Promise { - if (this.state.appState.streamingPhase !== 'idle') { - this.showError('Cannot undo while streaming — press Esc or Ctrl-C first.'); - return; - } - - const session = this.session; - if (session === undefined) { - this.showError(NO_ACTIVE_SESSION_MESSAGE); - return; - } - - const entries = this.state.transcriptEntries; - const lastUserIndex = entries.findLastIndex(isUndoAnchorEntry); - if (lastUserIndex < 0) { - this.showError('Nothing to undo.'); - return; - } - - try { - await session.undoHistory(1); - } catch (error) { - const message = formatErrorMessage(error); - this.showError(`Failed to undo: ${message}`); - return; - } - - const children = this.state.transcriptContainer.children; - let lastUserComponentIndex = -1; - for (let i = children.length - 1; i >= 0; i--) { - const child = children[i]; - if ( - child instanceof UserMessageComponent || - (child instanceof SkillActivationComponent && child.trigger === 'user-slash') - ) { - lastUserComponentIndex = i; - break; - } - } - - if (lastUserComponentIndex >= 0) { - children.splice(lastUserComponentIndex); - this.state.transcriptContainer.invalidate(); - } - - entries.splice(lastUserIndex); - - if (entries.length === 0) { - this.renderWelcome(); - } - - this.state.ui.requestRender(); - } - async getStartupMcpMs(): Promise { const session = this.session; if (session === undefined) return 0; @@ -1835,10 +1781,3 @@ export class KimiTUI { } } - -function isUndoAnchorEntry(entry: TranscriptEntry): boolean { - return ( - entry.kind === 'user' || - (entry.kind === 'skill_activation' && entry.skillTrigger === 'user-slash') - ); -} diff --git a/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts b/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts index 42bf0ca9..59035e43 100644 --- a/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts +++ b/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts @@ -57,7 +57,6 @@ interface MessageDriver { handleUserInput(text: string): void; persistInputHistory(text: string): Promise; getCurrentSessionId(): string; - undoLastTurn(): Promise; } interface FeedbackDriver extends MessageDriver { @@ -699,9 +698,17 @@ describe('KimiTUI message flow', () => { driver.handleUserInput('hello'); driver.state.appState.streamingPhase = 'idle'; - await driver.undoLastTurn(); + driver.handleUserInput('/undo'); + + await vi.waitFor(() => { + expect(session.undoHistory).toHaveBeenCalledWith(1); + }); + await vi.waitFor(() => { + expect(stripSgr(renderTranscript(driver))).toContain( + 'Error: Failed to undo: core rpc unavailable', + ); + }); - expect(session.undoHistory).toHaveBeenCalledWith(1); expect(driver.state.transcriptEntries).toEqual([ expect.objectContaining({ kind: 'user', @@ -710,7 +717,6 @@ describe('KimiTUI message flow', () => { ]); const transcript = stripSgr(renderTranscript(driver)); expect(transcript).toContain('hello'); - expect(transcript).toContain('Error: Failed to undo: core rpc unavailable'); }); it('does not duplicate welcome after undoing the only turn', async () => { @@ -719,7 +725,11 @@ describe('KimiTUI message flow', () => { driver.handleUserInput('hello'); driver.state.appState.streamingPhase = 'idle'; - await driver.undoLastTurn(); + driver.handleUserInput('/undo'); + + await vi.waitFor(() => { + expect(driver.state.transcriptEntries).toEqual([]); + }); expect( driver.state.transcriptContainer.children.filter( @@ -744,7 +754,11 @@ describe('KimiTUI message flow', () => { ); driver.state.appState.streamingPhase = 'idle'; - await driver.undoLastTurn(); + driver.handleUserInput('/undo'); + + await vi.waitFor(() => { + expect(driver.state.transcriptEntries).toEqual([]); + }); expect(driver.state.transcriptEntries).toEqual([]); const transcript = stripSgr(renderTranscript(driver)); @@ -768,7 +782,16 @@ describe('KimiTUI message flow', () => { ); driver.state.appState.streamingPhase = 'idle'; - await driver.undoLastTurn(); + driver.handleUserInput('/undo'); + + await vi.waitFor(() => { + expect(driver.state.transcriptEntries).toEqual([ + expect.objectContaining({ + kind: 'user', + content: 'hello', + }), + ]); + }); expect(driver.state.transcriptEntries).toEqual([ expect.objectContaining({ From f99dca96484a70c2491f6882c949546c7b4b39e5 Mon Sep 17 00:00:00 2001 From: _Kerman Date: Tue, 2 Jun 2026 11:31:12 +0800 Subject: [PATCH 08/10] fix --- apps/kimi-code/src/tui/commands/undo.ts | 57 ++++++++++++++++++- .../test/tui/kimi-tui-message-flow.test.ts | 23 ++++++++ 2 files changed, 78 insertions(+), 2 deletions(-) diff --git a/apps/kimi-code/src/tui/commands/undo.ts b/apps/kimi-code/src/tui/commands/undo.ts index 46dbb6c3..4e7ab0a6 100644 --- a/apps/kimi-code/src/tui/commands/undo.ts +++ b/apps/kimi-code/src/tui/commands/undo.ts @@ -1,5 +1,12 @@ import { WelcomeComponent } from '../components/chrome/welcome'; +import { AgentGroupComponent } from '../components/messages/agent-group'; +import { AssistantMessageComponent } from '../components/messages/assistant-message'; +import { BackgroundAgentStatusComponent } from '../components/messages/background-agent-status'; +import { CronMessageComponent } from '../components/messages/cron-message'; +import { ReadGroupComponent } from '../components/messages/read-group'; import { SkillActivationComponent } from '../components/messages/skill-activation'; +import { ThinkingComponent } from '../components/messages/thinking'; +import { ToolCallComponent } from '../components/messages/tool-call'; import { UserMessageComponent } from '../components/messages/user-message'; import { NO_ACTIVE_SESSION_MESSAGE } from '../constant/kimi-tui'; import type { TranscriptEntry } from '../types'; @@ -51,11 +58,14 @@ export async function handleUndoCommand(host: SlashCommandHost): Promise { } if (lastUserComponentIndex >= 0) { - children.splice(lastUserComponentIndex); + removeUndoContextComponents(children, lastUserComponentIndex); host.state.transcriptContainer.invalidate(); } - entries.splice(lastUserIndex); + const preservedEntries = entries.slice(lastUserIndex).filter( + (entry) => !isUndoContextEntry(entry), + ); + entries.splice(lastUserIndex, entries.length - lastUserIndex, ...preservedEntries); if (entries.length === 0) { renderWelcome(host); @@ -71,6 +81,49 @@ function isUndoAnchorEntry(entry: TranscriptEntry): boolean { ); } +function isUndoContextEntry(entry: TranscriptEntry): boolean { + switch (entry.kind) { + case 'user': + case 'assistant': + case 'tool_call': + case 'thinking': + case 'skill_activation': + case 'cron': + return true; + case 'status': + case 'welcome': + return false; + } +} + +function removeUndoContextComponents( + children: SlashCommandHost['state']['transcriptContainer']['children'], + startIndex: number, +): void { + for (let i = children.length - 1; i >= startIndex; i--) { + const child = children[i]; + if (child !== undefined && isUndoContextComponent(child)) { + children.splice(i, 1); + } + } +} + +function isUndoContextComponent( + child: SlashCommandHost['state']['transcriptContainer']['children'][number], +): boolean { + return ( + child instanceof UserMessageComponent || + child instanceof AssistantMessageComponent || + child instanceof ThinkingComponent || + child instanceof ToolCallComponent || + child instanceof AgentGroupComponent || + child instanceof ReadGroupComponent || + child instanceof SkillActivationComponent || + child instanceof BackgroundAgentStatusComponent || + child instanceof CronMessageComponent + ); +} + function renderWelcome(host: SlashCommandHost): void { if ( host.state.transcriptContainer.children.some( diff --git a/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts b/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts index 59035e43..c96f0611 100644 --- a/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts +++ b/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts @@ -738,6 +738,29 @@ describe('KimiTUI message flow', () => { ).toHaveLength(1); }); + it('keeps command notices that are not part of the undone context', async () => { + const { driver, session } = await makeDriver(); + + driver.handleUserInput('hello'); + driver.state.appState.streamingPhase = 'idle'; + driver.handleUserInput('/auto on'); + + await vi.waitFor(() => { + expect(stripSgr(renderTranscript(driver))).toContain('Auto mode: ON'); + }); + + driver.handleUserInput('/undo'); + + await vi.waitFor(() => { + expect(session.undoHistory).toHaveBeenCalledWith(1); + }); + + const transcript = stripSgr(renderTranscript(driver)); + expect(transcript).not.toContain('hello'); + expect(transcript).toContain('Auto mode: ON'); + expect(driver.state.appState.permissionMode).toBe('auto'); + }); + it('undoes from the real user turn when the last skill activation came from the model', async () => { const { driver } = await makeDriver(); From 2ebc40722c432eba62fe79fb9e3c695ff131319f Mon Sep 17 00:00:00 2001 From: _Kerman Date: Tue, 2 Jun 2026 11:36:38 +0800 Subject: [PATCH 09/10] fix --- apps/kimi-code/src/tui/commands/dispatch.ts | 2 +- apps/kimi-code/src/tui/commands/undo.ts | 85 ++++++++++++++----- .../test/tui/kimi-tui-message-flow.test.ts | 51 +++++++++++ 3 files changed, 116 insertions(+), 22 deletions(-) diff --git a/apps/kimi-code/src/tui/commands/dispatch.ts b/apps/kimi-code/src/tui/commands/dispatch.ts index 993aeb61..930dcbb9 100644 --- a/apps/kimi-code/src/tui/commands/dispatch.ts +++ b/apps/kimi-code/src/tui/commands/dispatch.ts @@ -282,7 +282,7 @@ async function handleBuiltInSlashCommand( await handleLogoutCommand(host); return; case 'undo': - await handleUndoCommand(host); + await handleUndoCommand(host, args); return; default: host.showError(`Unknown slash command: /${String(name)}`); diff --git a/apps/kimi-code/src/tui/commands/undo.ts b/apps/kimi-code/src/tui/commands/undo.ts index 4e7ab0a6..578547c8 100644 --- a/apps/kimi-code/src/tui/commands/undo.ts +++ b/apps/kimi-code/src/tui/commands/undo.ts @@ -1,3 +1,5 @@ +import type { Component } from '@earendil-works/pi-tui'; + import { WelcomeComponent } from '../components/chrome/welcome'; import { AgentGroupComponent } from '../components/messages/agent-group'; import { AssistantMessageComponent } from '../components/messages/assistant-message'; @@ -17,12 +19,21 @@ import type { SlashCommandHost } from './dispatch'; // Undo command // --------------------------------------------------------------------------- -export async function handleUndoCommand(host: SlashCommandHost): Promise { +export async function handleUndoCommand( + host: SlashCommandHost, + args: string = '', +): Promise { if (host.state.appState.streamingPhase !== 'idle') { host.showError('Cannot undo while streaming — press Esc or Ctrl-C first.'); return; } + const count = parseUndoCount(args); + if (count === undefined) { + host.showError('Usage: /undo [count], where count is a positive integer.'); + return; + } + const session = host.session; if (session === undefined) { host.showError(NO_ACTIVE_SESSION_MESSAGE); @@ -30,14 +41,14 @@ export async function handleUndoCommand(host: SlashCommandHost): Promise { } const entries = host.state.transcriptEntries; - const lastUserIndex = entries.findLastIndex(isUndoAnchorEntry); - if (lastUserIndex < 0) { + const lastUserIndex = findUndoAnchorEntryIndex(entries, count); + if (lastUserIndex === undefined) { host.showError('Nothing to undo.'); return; } try { - await session.undoHistory(1); + await session.undoHistory(count); } catch (error) { const message = formatErrorMessage(error); host.showError(`Failed to undo: ${message}`); @@ -45,19 +56,8 @@ export async function handleUndoCommand(host: SlashCommandHost): Promise { } const children = host.state.transcriptContainer.children; - let lastUserComponentIndex = -1; - for (let i = children.length - 1; i >= 0; i--) { - const child = children[i]; - if ( - child instanceof UserMessageComponent || - (child instanceof SkillActivationComponent && child.trigger === 'user-slash') - ) { - lastUserComponentIndex = i; - break; - } - } - - if (lastUserComponentIndex >= 0) { + const lastUserComponentIndex = findUndoAnchorComponentIndex(children, count); + if (lastUserComponentIndex !== undefined) { removeUndoContextComponents(children, lastUserComponentIndex); host.state.transcriptContainer.invalidate(); } @@ -74,6 +74,14 @@ export async function handleUndoCommand(host: SlashCommandHost): Promise { host.state.ui.requestRender(); } +function parseUndoCount(args: string): number | undefined { + const value = args.trim(); + if (value.length === 0) return 1; + if (!/^[1-9]\d*$/.test(value)) return undefined; + const count = Number(value); + return Number.isSafeInteger(count) ? count : undefined; +} + function isUndoAnchorEntry(entry: TranscriptEntry): boolean { return ( entry.kind === 'user' || @@ -81,6 +89,21 @@ function isUndoAnchorEntry(entry: TranscriptEntry): boolean { ); } +function findUndoAnchorEntryIndex( + entries: readonly TranscriptEntry[], + count: number, +): number | undefined { + let found = 0; + for (let i = entries.length - 1; i >= 0; i--) { + const entry = entries[i]; + if (entry !== undefined && isUndoAnchorEntry(entry)) { + found++; + if (found === count) return i; + } + } + return undefined; +} + function isUndoContextEntry(entry: TranscriptEntry): boolean { switch (entry.kind) { case 'user': @@ -96,8 +119,23 @@ function isUndoContextEntry(entry: TranscriptEntry): boolean { } } +function findUndoAnchorComponentIndex( + children: readonly Component[], + count: number, +): number | undefined { + let found = 0; + for (let i = children.length - 1; i >= 0; i--) { + const child = children[i]; + if (child !== undefined && isUndoAnchorComponent(child)) { + found++; + if (found === count) return i; + } + } + return undefined; +} + function removeUndoContextComponents( - children: SlashCommandHost['state']['transcriptContainer']['children'], + children: Component[], startIndex: number, ): void { for (let i = children.length - 1; i >= startIndex; i--) { @@ -108,9 +146,14 @@ function removeUndoContextComponents( } } -function isUndoContextComponent( - child: SlashCommandHost['state']['transcriptContainer']['children'][number], -): boolean { +function isUndoAnchorComponent(child: Component): boolean { + return ( + child instanceof UserMessageComponent || + (child instanceof SkillActivationComponent && child.trigger === 'user-slash') + ); +} + +function isUndoContextComponent(child: Component): boolean { return ( child instanceof UserMessageComponent || child instanceof AssistantMessageComponent || diff --git a/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts b/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts index c96f0611..d2acfdf4 100644 --- a/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts +++ b/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts @@ -761,6 +761,57 @@ describe('KimiTUI message flow', () => { expect(driver.state.appState.permissionMode).toBe('auto'); }); + it('undoes multiple turns when a count is provided', async () => { + const { driver, session } = await makeDriver(); + + driver.handleUserInput('first'); + driver.state.appState.streamingPhase = 'idle'; + driver.handleUserInput('second'); + driver.state.appState.streamingPhase = 'idle'; + driver.handleUserInput('third'); + driver.state.appState.streamingPhase = 'idle'; + + driver.handleUserInput('/undo 2'); + + await vi.waitFor(() => { + expect(session.undoHistory).toHaveBeenCalledWith(2); + }); + + expect(driver.state.transcriptEntries).toEqual([ + expect.objectContaining({ + kind: 'user', + content: 'first', + }), + ]); + const transcript = stripSgr(renderTranscript(driver)); + expect(transcript).toContain('first'); + expect(transcript).not.toContain('second'); + expect(transcript).not.toContain('third'); + }); + + it('rejects invalid undo counts without changing context', async () => { + const { driver, session } = await makeDriver(); + + driver.handleUserInput('hello'); + driver.state.appState.streamingPhase = 'idle'; + + driver.handleUserInput('/undo 0'); + + await vi.waitFor(() => { + expect(stripSgr(renderTranscript(driver))).toContain( + 'Error: Usage: /undo [count], where count is a positive integer.', + ); + }); + + expect(session.undoHistory).not.toHaveBeenCalled(); + expect(driver.state.transcriptEntries).toEqual([ + expect.objectContaining({ + kind: 'user', + content: 'hello', + }), + ]); + }); + it('undoes from the real user turn when the last skill activation came from the model', async () => { const { driver } = await makeDriver(); From fd6bd63b041768df3ac1d231f8933eae616db310 Mon Sep 17 00:00:00 2001 From: _Kerman Date: Tue, 2 Jun 2026 12:34:56 +0800 Subject: [PATCH 10/10] fix --- apps/kimi-code/src/tui/commands/undo.ts | 7 ++ apps/kimi-code/src/tui/kimi-tui.ts | 3 + .../utils/transcript-component-metadata.ts | 15 +++ .../test/tui/kimi-tui-message-flow.test.ts | 91 +++++++++++++++++++ 4 files changed, 116 insertions(+) create mode 100644 apps/kimi-code/src/tui/utils/transcript-component-metadata.ts diff --git a/apps/kimi-code/src/tui/commands/undo.ts b/apps/kimi-code/src/tui/commands/undo.ts index 578547c8..752fd27a 100644 --- a/apps/kimi-code/src/tui/commands/undo.ts +++ b/apps/kimi-code/src/tui/commands/undo.ts @@ -13,6 +13,7 @@ import { UserMessageComponent } from '../components/messages/user-message'; import { NO_ACTIVE_SESSION_MESSAGE } from '../constant/kimi-tui'; import type { TranscriptEntry } from '../types'; import { formatErrorMessage } from '../utils/event-payload'; +import { getTranscriptComponentEntry } from '../utils/transcript-component-metadata'; import type { SlashCommandHost } from './dispatch'; // --------------------------------------------------------------------------- @@ -114,6 +115,7 @@ function isUndoContextEntry(entry: TranscriptEntry): boolean { case 'cron': return true; case 'status': + return entry.turnId !== undefined; case 'welcome': return false; } @@ -154,6 +156,11 @@ function isUndoAnchorComponent(child: Component): boolean { } function isUndoContextComponent(child: Component): boolean { + const entry = getTranscriptComponentEntry(child); + if (entry !== undefined) { + return isUndoContextEntry(entry); + } + return ( child instanceof UserMessageComponent || child instanceof AssistantMessageComponent || diff --git a/apps/kimi-code/src/tui/kimi-tui.ts b/apps/kimi-code/src/tui/kimi-tui.ts index 66458ed0..589ef5a1 100644 --- a/apps/kimi-code/src/tui/kimi-tui.ts +++ b/apps/kimi-code/src/tui/kimi-tui.ts @@ -120,6 +120,7 @@ import { installTerminalFocusTracking } from './utils/terminal-focus'; import { notifyTerminalOnce } from './utils/terminal-notification'; import { installTerminalThemeTracking } from './utils/terminal-theme'; import { detectTmuxKeyboardWarning } from './utils/tmux-keyboard'; +import { markTranscriptComponent } from './utils/transcript-component-metadata'; import { nextTranscriptId } from './utils/transcript-id'; export type { TUIState } from './tui-state'; @@ -1242,6 +1243,7 @@ export class KimiTUI { this.state.transcriptEntries.push(entry); const component = this.createTranscriptComponent(entry); if (component) { + markTranscriptComponent(component, entry); this.state.transcriptContainer.addChild(component); this.state.ui.requestRender(); } @@ -1268,6 +1270,7 @@ export class KimiTUI { this.appendTranscriptEntry({ id: nextTranscriptId(), kind: 'status', + turnId: request.turnId === undefined ? undefined : String(request.turnId), renderMode: 'notice', content: parts.join(''), }); diff --git a/apps/kimi-code/src/tui/utils/transcript-component-metadata.ts b/apps/kimi-code/src/tui/utils/transcript-component-metadata.ts new file mode 100644 index 00000000..12151958 --- /dev/null +++ b/apps/kimi-code/src/tui/utils/transcript-component-metadata.ts @@ -0,0 +1,15 @@ +import type { Component } from '@earendil-works/pi-tui'; + +import type { TranscriptEntry } from '../types'; + +const componentEntries = new WeakMap(); + +export function markTranscriptComponent(component: Component, entry: TranscriptEntry): void { + componentEntries.set(component, entry); +} + +export function getTranscriptComponentEntry( + component: Component, +): TranscriptEntry | undefined { + return componentEntries.get(component); +} diff --git a/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts b/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts index d2acfdf4..f23e475c 100644 --- a/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts +++ b/apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts @@ -761,6 +761,97 @@ describe('KimiTUI message flow', () => { expect(driver.state.appState.permissionMode).toBe('auto'); }); + it('removes turn-scoped background status entries and restores welcome', async () => { + const { driver, session } = await makeDriver(); + + driver.handleUserInput('hello'); + driver.state.appState.streamingPhase = 'idle'; + driver.sessionEventHandler.handleEvent( + { + type: 'background.task.started', + agentId: 'main', + sessionId: 'ses-1', + turnId: 1, + info: { + taskId: 'bash-bg123456', + command: 'npm test', + description: 'Run tests in background', + status: 'running', + pid: 1234, + exitCode: null, + startedAt: Date.now(), + endedAt: null, + }, + } as Event, + () => {}, + ); + + await vi.waitFor(() => { + const transcript = stripSgr(renderTranscript(driver)); + expect(transcript).toContain('bash task started in background'); + expect(transcript).toContain('Run tests in background'); + }); + + driver.handleUserInput('/undo'); + + await vi.waitFor(() => { + expect(session.undoHistory).toHaveBeenCalledWith(1); + }); + + const transcript = stripSgr(renderTranscript(driver)); + expect(driver.state.transcriptEntries).toEqual([]); + expect(transcript).not.toContain('hello'); + expect(transcript).not.toContain('bash task started in background'); + expect(transcript).not.toContain('Run tests in background'); + expect( + driver.state.transcriptContainer.children.filter( + (child) => child instanceof WelcomeComponent, + ), + ).toHaveLength(1); + }); + + it('removes approval notices from undone turns', async () => { + const { driver, session } = await makeDriver(); + const approvalHandler = vi.mocked(session.setApprovalHandler).mock.calls[0]?.[0] as + | ((request: ApprovalRequest) => Promise) + | undefined; + if (approvalHandler === undefined) throw new Error('expected approval handler'); + + driver.handleUserInput('hello'); + driver.state.appState.streamingPhase = 'idle'; + const response = approvalHandler({ + turnId: 1, + toolCallId: 'call_bash', + toolName: 'Bash', + action: 'Run shell command', + display: { + kind: 'generic', + summary: 'Run shell command', + detail: { command: 'echo ok', description: 'Run a shell command' }, + }, + }); + + await vi.waitFor(() => { + expect(driver.state.editorContainer.children[0]).toBeInstanceOf(ApprovalPanelComponent); + }); + (driver.state.editorContainer.children[0] as ApprovalPanelComponent).handleInput('1'); + await expect(response).resolves.toMatchObject({ decision: 'approved' }); + + await vi.waitFor(() => { + expect(stripSgr(renderTranscript(driver))).toContain('Approved: Run shell command'); + }); + + driver.handleUserInput('/undo'); + + await vi.waitFor(() => { + expect(session.undoHistory).toHaveBeenCalledWith(1); + }); + + const transcript = stripSgr(renderTranscript(driver)); + expect(transcript).not.toContain('hello'); + expect(transcript).not.toContain('Approved: Run shell command'); + }); + it('undoes multiple turns when a count is provided', async () => { const { driver, session } = await makeDriver();